PyTorch - Convert CIFAR dataset to `TensorDataset`

Question

I train ResNet34 on CIFAR dataset. For a certain reason, I need to convert the dataset into TensorDataset. My solution is based on this: https://stackoverflow.com/a