Set `find_unused_parameters` to True
@@ -284,7 +284,9 @@ class Trainer(TrainerBase):
# normal case, reference to `Model` is retained. This reference is only used in
# these places: `model.__call__`, `model.train` and `model.eval`.
if self._distributed:
- self._pytorch_model = DistributedDataParallel(self.model, device_ids=[self.cuda_device])
+ self._pytorch_model = DistributedDataParallel(
+ self.model, device_ids=[self.cuda_device], find_unused_parameters=True
+ )
else:
self._pytorch_model = self.model
