Lightning-AI#515 Add on_gpu

dreamgonfly · Nov 18, 2019 · 3477108 · 3477108
1 parent 70cbd8c
commit 3477108
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py
@@ -46,7 +46,7 @@ def restore_weights(self, model):
 
  if not did_restore_hpc_weights:
  if self.resume_from_checkpoint is not None:
- self.restore(self.resume_from_checkpoint)
+ self.restore(self.resume_from_checkpoint, on_gpu=self.on_gpu)
  else:
  # restore weights if same exp version
  self.restore_state_if_checkpoint_exists(model)