Non plain-tensor batches
Hello!
I have been trying to use Lightning to train a graph neural network built with torch_geometric package using a GPU.
This is the error I get when I try to fit the model:
~/miniconda3/envs/pyg/lib/python3.7/site-packages/torch_geometric/nn/conv/gcn_conv.py in forward(self, x, edge_index, edge_weight)
     83     def forward(self, x, edge_index, edge_weight=None):
     84         """"""
---> 85         x = torch.matmul(x, self.weight)
     86 
     87         if self.cached and self.cached_result is not None:

RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 'mat2'
I investigated the code of Lightning for a probable cause and found that transfer_batch_to_gpu causes this error.
The current behavior of this function considers a batch to be either a 'plain' torch.Tensor or some simple collection of such objects (a list, a dict or a tuple). The problem is that torch_geometric uses a custom aggregate type Data (docs) which itself implements to method (and does not implement cuda, I believe).
It would be nice if you could make the code more flexible to process this case (and similar cases in general) correctly. I believe, the best solution is to replace isinstance(batch, torch.Tensor) by a condition of batch object having either to or cuda method implemented. This should solve the problem in general.