From 60d4b803224d79a0c8f536302b8e4dbebccedde1 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Thu, 16 May 2019 15:55:21 -0400 Subject: [PATCH] added amp level option --- pytorch_lightning/models/trainer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index d360d7ac..78689f52 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -372,8 +372,9 @@ class Trainer(TrainerIO): if self.use_amp: for optimizer in self.optimizers: with amp.scale_loss(loss, optimizer) as scaled_loss: - optimizer.backward(scaled_loss) - # scaled_loss.backward() + scaled_loss.backward() + for param in self.model.parameters(): + print(param.grad.float().sum()) else: loss.backward()