Skip to content

Commit 3880893

Browse files
authored
Fix neural style transfer scripmodule creation script error (#48)
1 parent 2bc1d05 commit 3880893

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

tutorials/advanced/neural_style_transfer/model/create_vgg19_layers_scriptmodule.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,13 @@ def main():
99
for param in vgg_19_layers.parameters():
1010
param.requires_grad = False
1111

12+
example = torch.rand(1, 3, 224, 224)
13+
14+
traced_script_module = torch.jit.trace(vgg_19_layers, example)
15+
1216
# Serialize scriptmodule to a file.
1317
filename = "vgg19_layers.pt"
14-
vgg_19_layers.save(filename)
18+
traced_script_module.save(filename)
1519
print(f"Successfully created scriptmodule file {filename}.")
1620

1721

0 commit comments

Comments
 (0)