diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index b9000ff4..acff43fc 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -319,6 +319,7 @@ class Trainer(TrainerIO): elif self.use_dp: self.dp_train(model) + # ON CPU else: # CHOOSE OPTIMIZER # filter out the weights that were done on gpu so we can load on good old cpus