From a2b20b46bca5101627ed392aec17611ac0e97133 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ayberk=20Ayd=C4=B1n?= Date: Tue, 21 Jan 2020 16:09:27 +0300 Subject: [PATCH] remove unnecesarry gradient freeze/unfreeze for single optimizer setup (#719) --- pytorch_lightning/trainer/training_loop.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 48a10e19..06d3ebd5 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -460,11 +460,12 @@ class TrainerTrainLoopMixin(ABC): for opt_idx, optimizer in enumerate(self.optimizers): # make sure only the gradients of the current optimizer's paramaters are calculated # in the training step to prevent dangling gradients in multiple-optimizer setup. - for param in self.get_model().parameters(): - param.requires_grad = False - for group in optimizer.param_groups: - for param in group['params']: - param.requires_grad = True + if len(self.optimizers) > 1: + for param in self.get_model().parameters(): + param.requires_grad = False + for group in optimizer.param_groups: + for param in group['params']: + param.requires_grad = True # wrap the forward step in a closure so second order methods work def optimizer_closure():