question: How to implement channels last memory format callback
context: Hi there
  Pytorch docs recommends using channels last when training vision models in mixed precision. To enable, you need to do two changes:
  
  Move you model to channel last format: model = model.to(memory_format=torch.channels_last) # Replace with your model. This can be done in on_fit_start callback hook as model.to performs an in place modification:
  
  class ChannelsLast(pl.Callback):
      def on_fit_start(self, trainer, pl_module: pl.LightningModule) -> None:
          # Inplace model modification
          pl_module.to(memory_format=torch.channels_last)
  
  
  Move input data to channel last format before feeding to the model: input = input.to(memory_format=torch.channels_last).
  
  My problem is in step 2. I don't find any PyTorch lightning hook that allows me to make this modification to the batch :/. The only options left are to add it as data transforms (that must be used in conjunction with the callback) or doing all channel last related logic inside the LightningModule. I would prefer to avoid this last solution as it could clutter the LightningModule with unnecessary code.
  Do you know a to do step 2 in a callback?
answer: You can use any of the following if done inside the LightningModule:
  batch = self.on_before_batch_transfer(batch, dataloader_idx)
  batch = self.transfer_batch_to_device(batch, device)
  batch = self.on_after_batch_transfer(batch, dataloader_idx)
  If you really need to do it in the Callback, I guess you could use on_train_batch_start since the modification is in-place. But I wouldn't recommend it.