changed to prepare_batch_for_transfer

Author: rohitgr7

pytorch_lightning/core/decorators.py

@@ -57,8 +57,8 @@ def auto_transfer_args(self, *args, **kwargs):
         if not isinstance(self, LightningModule):
             return fn(self, *args, **kwargs)
 
-        args = self.transfer_batch_to_device(args, self.device)
-        kwargs = self.transfer_batch_to_device(kwargs, self.device)
+        args = self.prepare_batch_for_transfer(args, self.device)
+        kwargs = self.prepare_batch_for_transfer(kwargs, self.device)
         return fn(self, *args, **kwargs)
 
     return auto_transfer_args

pytorch_lightning/core/memory.py

@@ -219,7 +219,7 @@ def _forward_example_input(self) -> None:
         trainer = self._model.trainer
 
         input_ = model.example_input_array
-        input_ = model.transfer_batch_to_device(input_, model.device)
+        input_ = model.prepare_batch_for_transfer(input_, model.device)
 
         if trainer is not None and trainer.amp_backend == AMPType.NATIVE and not trainer.use_tpu:
             model.forward = torch.cuda.amp.autocast()(model.forward)

pytorch_lightning/loggers/tensorboard.py

@@ -197,7 +197,7 @@ def log_graph(self, model: LightningModule, input_array=None):
                 input_array = model.example_input_array
 
             if input_array is not None:
-                input_array = model.transfer_batch_to_device(input_array, model.device)
+                input_array = model.prepare_batch_for_transfer(input_array, model.device)
                 self.experiment.add_graph(model, input_array)
             else:
                 rank_zero_warn('Could not log computational graph since the'

pytorch_lightning/loggers/test_tube.py

@@ -146,7 +146,7 @@ def log_graph(self, model: LightningModule, input_array=None):
             if input_array is not None:
                 self.experiment.add_graph(
                     model,
-                    model.transfer_batch_to_device(
+                    model.prepare_batch_for_transfer(
                         model.example_input_array, model.device)
                 )
             else:

pytorch_lightning/trainer/connectors/data_connector.py

@@ -117,7 +117,8 @@ def attach_datamodule(self, model, datamodule, stage):
             if is_overridden('test_dataloader', datamodule):
                 model.test_dataloader = datamodule.test_dataloader
 
-            # Override transfer_batch_to_device if dataset-specific to_device logic has been defined in datamodule
+            # TODO check this before merge
+            # Override prepare_batch_for_transfer if dataset-specific to_device logic has been defined in datamodule
             model.prepare_batch_for_transfer = datamodule.prepare_batch_for_transfer
 
             # TODO remove this after all the changes are done