Make VisionDataset generic following Dataset

[Glinte]

🚀 The feature

torch.utils.data.datasets.Dataset is already generic wrt the item type, so it would be nice to expose the type parameter to users

_T_co = TypeVar("_T_co", covariant=True)

class Dataset(Generic[_T_co]):
    def __getitem__(self, index) -> _T_co:
        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")

but

class VisionDataset(data.Dataset):
    ...

Motivation, pitch

I am making a subclass of VisionDataset and I noticed DataLoader(VisionDataset(...)) doesn't have generic information.

Alternatives

No response

Additional context

No response