Allow passing dict object as batch
🚀 Feature
Allow the batch variable to be a (nested) dict object (in training_step(), etc.)
Motivation
I was porting my code from my own implementation of a Trainer to Pytorch Lightning, as I believe this is an asset for reproducibility and clarity.
However I stumbled across a problem, it is not possible for the variable batch to be a (nested) dict object and it can only be a Tensor, tuple or list.
Pitch
The dataset structure that I am using is the following:
import torch
from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(
        self, n: int
    ):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, item):
        frame = {
            "index": item,
            "A": {
                "image": torch.randn(3, 128, 128),
                "R": torch.randn(3, 3),
                "t": torch.randn(3, 1),
                "K": torch.randn(3, 3),
            },
            "B": {
                "image": torch.randn(3, 128, 128),
                "R": torch.randn(3, 3),
                "t": torch.randn(3, 1),
                "K": torch.randn(3, 3),
            },
        }
        return frame
The PyTorch Dataloader is able to make a batch of frames with the default collate_fn.
To move things onto the GPU, I am currently doing:
def move(self, d: dict, device) -> dict:
    for k in d:
        if isinstance(d[k], dict):
            d[k] = self.move(d[k])
        elif isinstance(d[k], (Tensor)):
            d[k] = d[k].to(device=device, non_blocking=True)
    return d
There is probably a nicer way to do it, but it does the job.
Could you please add this feature, as passing a tuple is very impractical when the amount of variables is high?
Thanks in advance,
Guillaume