How to implement channels last memory format callback
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?
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.