diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index bb63e35b..3daa759d 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -283,7 +283,6 @@ class Trainer(TrainerIO): model, self.optimizers, opt_level=self.amp_level, ) self.optimizers = optimizers - model.trainer = self # when using gpus, first thing we do is spawn a new process between each worker # applies to single gpu, multi-gpu and multi-nodes @@ -385,6 +384,8 @@ class Trainer(TrainerIO): if self.on_gpu: ref_model = model.module + ref_model.trainer = self + # set local properties on the model ref_model.on_gpu = self.on_gpu