Training loop takes a long time each epoch using TensorDataset

[ptrblck]

If you are dealing with a data loading bottleneck, I would recommend to read this post which gives a good overview about possible reasons and workarounds.