release v0.3.4.1

pytorch_lightning/models/trainer.py

@@ -178,9 +178,9 @@ class Trainer(TrainerIO):
             self.use_ddp = distributed_backend == 'ddp'
 
             # use ddp automatically if nb_gpu_nodes > 1
-            if nb_gpu_nodes > 1:
+            if nb_gpu_nodes > 1 and self.use_dp:
                 self.use_ddp = True
-                self.use_ddp = False
+                self.use_dp = False
                 w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
                     'Switching to DistributedDataParallel for you. ' \
                     'To silence this warning set distributed_backend=ddp'
@@ -443,7 +443,6 @@ class Trainer(TrainerIO):
         self.optimizers = model.configure_optimizers()
 
         model.cuda(self.data_parallel_device_ids[0])
-        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
 
         # run through amp wrapper
         if self.use_amp:
@@ -453,6 +452,8 @@ class Trainer(TrainerIO):
             )
             self.optimizers = optimizers
 
+        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
         self.__run_pretrain_routine(model)
 
     def ddp_train(self, gpu_nb, model):
@@ -506,7 +507,7 @@ class Trainer(TrainerIO):
             )
             self.optimizers = optimizers
 
-        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb], find_unused_parameters=True)
 
         # continue training routine
         self.__run_pretrain_routine(model)
@@ -775,6 +776,11 @@ class Trainer(TrainerIO):
         else:
             loss.backward()
 
+        # insert after step hook
+        if self.__is_function_implemented('on_after_backward'):
+            model_ref = self.__get_model()
+            response = model_ref.on_after_backward()
+
         if self.print_nan_grads:
             model = self.__get_model()
             for param in model.parameters():

pytorch_lightning/root_module/hooks.py

@@ -36,3 +36,10 @@ class ModelHooks(torch.nn.Module):
         """
         pass
 
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass
+

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.3.3',
+    version='0.3.4.1',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",