Fix jit doc model loading example

pytorch · Nov 17, 2020 · 6199e42 · 6199e42
1 parent 49f0e5d
commit 6199e42
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docs/source/jit.rst b/docs/source/jit.rst
@@ -547,10 +547,10 @@ best practices?
 
       cpu_model = gpu_model.cpu()
       sample_input_cpu = sample_input_gpu.cpu()
-      traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
+      traced_cpu = torch.jit.trace(cpu_model, sample_input_cpu)
       torch.jit.save(traced_cpu, "cpu.pth")
 
-      traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
+      traced_gpu = torch.jit.trace(gpu_model, sample_input_gpu)
       torch.jit.save(traced_gpu, "gpu.pth")
 
       # ... later, when using the model: