fix

Author: Borda

pytorch_lightning/accelerators/cpu_accelerator.py

@@ -61,22 +61,22 @@ def train(self):
         results = self.train_or_test()
         return results
 
-    def _step(self, model_step: Callable, args):
+    def _step(self, model_step: Callable, *args):
         if self.trainer.amp_backend == AMPType.NATIVE:
             with torch.cuda.amp.autocast():
                 output = model_step(*args)
         else:
             output = model_step(*args)
         return output
 
-    def training_step(self, args):
-        return self._step(self, self.trainer.model.training_step, args)
+    def training_step(self, *args):
+        return self._step(self.trainer.model.training_step, *args)
 
-    def validation_step(self, args):
-        return self._step(self, self.trainer.model.validation_step, args)
+    def validation_step(self, *args):
+        return self._step(self.trainer.model.validation_step, *args)
 
-    def test_step(self, args):
-        return self._step(self, self.trainer.model.test_step, args)
+    def test_step(self, *args):
+        return self._step(self.trainer.model.test_step, *args)
 
     def sync_tensor(self,
                     tensor: Union[torch.Tensor],

pytorch_lightning/accelerators/horovod_accelerator.py

@@ -114,7 +114,7 @@ def train(self):
         hvd.join()
         return results
 
-    def _step(self, model_step: Callable, args):
+    def _step(self, model_step: Callable, *args):
         if self.trainer.on_gpu:
             batch = args[0]
             batch = self.batch_to_device(batch, hvd.local_rank())
@@ -128,14 +128,14 @@ def _step(self, model_step: Callable, args):
 
         return output
 
-    def training_step(self, args):
-        return self._step(self, self.trainer.model.training_step, args)
+    def training_step(self, *args):
+        return self._step(self.trainer.model.training_step, *args)
 
-    def validation_step(self, args):
-        return self._step(self, self.trainer.model.validation_step, args)
+    def validation_step(self, *args):
+        return self._step(self.trainer.model.validation_step, *args)
 
-    def test_step(self, args):
-        return self._step(self, self.trainer.model.test_step, args)
+    def test_step(self, *args):
+        return self._step(self.trainer.model.test_step, *args)
 
     def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
         super().backward(closure_loss, optimizer, opt_idx, *args, **kwargs)