@auto_move_data unexpectedly uses transfer_batch_to_device from DataModule
🐛 Bug
I have a LightningDataModule that produces a custom batch object, and I have implemented transfer_batch_to_device to move this data to the GPU. This works.
I have a separate infer method (on my LightningModule) which is invoked with a Tensor, and I wanted to use @auto_move_data to move this Tensor to
the GPU. However, this unexpectedly uses my data module's transfer_batch_to_device method, and assumed that the input to infer will be of my custom batch type.
To Reproduce
https://colab.research.google.com/drive/18f8886hG9AVSkYdVi96CIcZjrl9zX6Np#scrollTo=AAtq1hwSmjKe
Expected behavior

I think this API is a bit confusing/surprising. I was expecting @auto_move_data to behave in the "default" mode since I passed in a Tensor, and I was not expecting it to use my  data module's transfer_batch_to_device method - especially since the infer method and the @auto_move_data decorator were on the LightningModule, not the data module.

IMHO, a less confusing API would be to expect a to() method on the custom batch object, and make transfer_batch_to_device just call that to() method. Then it would work uniformly for tensors and custom batch objects.

Another thing that makes this problematic: the signature of transfer_batch_to_device is supposed to be def transfer_batch_to_device(self, batch, device): https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html?highlight=transfer_batch_to_device#transfer-batch-to-device

But in @auto_move_data, it is getting called with a tuple of (args, kwargs), and no device AFAICT.
Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !

CUDA:

GPU:

Tesla T4


available:         True
version:           10.1


Packages:

numpy:             1.19.5
pyTorch_debug:     True
pyTorch_version:   1.7.0+cu101
pytorch-lightning: 1.2.0
tqdm:              4.41.1


System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.6.9
version:           #1 SMP Thu Jul 23 08:00:38 PDT 2020



Additional context