Thank you! This seems to have solved my main issue, but using the above code made the loader very slow. Actually, if I use it inside a loop:
for epoch in range(epochs):
       for x_ in train_loader:
           print(type(x_))

This takes five seconds to print the first time, and then I get an error message that the RAM is crushed for using all available memory!