make it private

Author: rohitgr7

pytorch_lightning/accelerators/accelerator.py

@@ -69,7 +69,7 @@ def train_or_test(self):
     def batch_to_device(self, batch: Any, device: torch.device):
         model = self.trainer.get_model()
         if model is not None:
-            return model.prepare_batch_for_transfer(batch, device)
+            return model._prepare_batch_for_transfer(batch, device)
         return move_data_to_device(batch, device)
 
     def training_step_end(self, output):

pytorch_lightning/core/hooks.py

@@ -560,15 +560,15 @@ def transfer_batch_to_device(self, batch, device):
         return move_data_to_device(batch, device)
 
     def on_before_batch_transfer(self, batch):
+        """
+        Called before batch is transfered to the device
+        """
         return batch
 
     def on_after_batch_transfer(self, batch):
-        return batch
-
-    def prepare_batch_for_transfer(self, batch: Any, device: Optional[torch.device] = None):
-        batch = self.on_before_batch_transfer(batch)
-        batch = self.transfer_batch_to_device(batch, device)
-        batch = self.on_after_batch_transfer(batch)
+        """
+        Called after batch is transfered to the device
+        """
         return batch
 
 

pytorch_lightning/core/lightning.py

@@ -171,6 +171,12 @@ def automatic_optimization(self) -> bool:
         """
         return True
 
+    def _prepare_batch_for_transfer(self, batch: Any, device: Optional[torch.device] = None):
+        batch = self.on_before_batch_transfer(batch)
+        batch = self.transfer_batch_to_device(batch, device)
+        batch = self.on_after_batch_transfer(batch)
+        return batch
+
     def print(self, *args, **kwargs) -> None:
         r"""
         Prints only from process 0. Use this in any distributed mode to log only once.
@@ -1787,7 +1793,7 @@ def to_torchscript(
                 if example_inputs is None:
                     example_inputs = self.example_input_array
                 # automatically send example inputs to the right device and use trace
-                example_inputs = self.prepare_batch_for_transfer(example_inputs)
+                example_inputs = self._prepare_batch_for_transfer(example_inputs)
                 torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
             else:
                 raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was:"

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.prepare_batch_for_transfer(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

@@ -198,7 +198,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.prepare_batch_for_transfer(input_array)
+                input_array = model._prepare_batch_for_transfer(input_array)
                 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):
                 input_array = model.example_input_array
 
             if input_array is not None:
-                input_array = self.prepare_batch_for_transfer(input_array)
+                input_array = self._prepare_batch_for_transfer(input_array)
                 self.experiment.add_graph(model, input_array)
             else:
                 rank_zero_warn('Could not log computational graph since the'