Which line of code threw this error?

To me the Subsets do not work well with TensorDataset, since Subsets are not the tensors. This line will fail:
train_ds = TensorDataset(train_ds, valid_ds)


Could you explain your confusion a bit?

I was not so sure why do we need Subsets. Since for MNIST dataset (60000 images) I cannot use them to make TensorDataset.
I am not sure if I can use something other than TensorDataset, but the idea was to
to move tensors on CUDA.
I guess it is faster to move data and targets on CUDA togeter, but I am not sure, maybe just data on CUDA is OK, and targes on CPU may be still be fast.
So based on making data and targets togeter I see TensorDataset as a good choice` for any unsupervised architecture, having data and targets.

Subset is a convenient way of using certain indices to specify subsets of your Dataset. While in your example you can easily split the loaded data manually, this won’t be possible in a lot of cases, where you would have to lazily load the data

This is maybe it, I haven’t experimented with lazy loading data.