[tune] missing torch.load in mnist_pytorch_trainable.py (#5103)

ray-project · Jul 6, 2019 · 4b56a5e · 4b56a5e
1 parent c5253cc
commit 4b56a5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/ray/tune/examples/mnist_pytorch_trainable.py b/python/ray/tune/examples/mnist_pytorch_trainable.py
@@ -167,7 +167,7 @@ def _save(self, checkpoint_dir):
         return checkpoint_path
 
     def _restore(self, checkpoint_path):
-        self.model.load_state_dict(checkpoint_path)
+        self.model.load_state_dict(torch.load(checkpoint_path))
 
 
 if __name__ == "__main__":