added amp level option

pytorch_lightning/models/trainer.py

@@ -33,6 +33,7 @@ class Trainer(TrainerIO):
                  log_save_interval=1, add_log_row_interval=1,
                  lr_scheduler_milestones=None,
                  use_amp=False,
+                 check_grad_nans=False,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
 
@@ -60,6 +61,7 @@ class Trainer(TrainerIO):
         self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
         self.lr_schedulers = []
         self.amp_level = amp_level
+        self.check_grad_nans = check_grad_nans
 
         # training state
         self.optimizers = None
@@ -373,11 +375,13 @@ class Trainer(TrainerIO):
             for optimizer in self.optimizers:
                 with amp.scale_loss(loss, optimizer) as scaled_loss:
                     scaled_loss.backward()
-                    for param in self.model.parameters():
-                        print(param.grad.float().sum())
         else:
             loss.backward()
 
+            if self.check_grad_nans:
+                for param in self.model.parameters():
+                    print(param.grad.float().sum())
+
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients

pytorch_lightning/utils/arg_parse.py

@@ -50,7 +50,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true')
     parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true')
     parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
-    parser.add_argument('--use_amp', dest='use_amp', action='store_true')
+    parser.add_argument('--use_amp', dest='use_amp', action='store_true')a=
+    parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')a=
     parser.add_argument('--amp_level', default='O2',type=str)