Dead loop when iterate through pytorch dataloader


I have a small CNN model that works fine on my PC-CPU.
I needed more computations power, so I moved all my data and code to Datalore platform.

Unfortunately, when I’m initializing my model over the Datalore Sheet, my model reaches only the lines where it needs to iterate over PyTorch dataloader and then entering into a dead loop.

This means that it can’t pull image samples for training, the model gets stuck, and I must restart the kernel.
If I’m using a batch size of 2, the model trains.
This bug does not occur when I’m using my own PC nor Kaggle platform with the same code for training.
I also tried to iterate over the dataloader outside of the model, using the code:

images, labels = iter(train_dataloader).next()

and this line of code reproduces the bug described above.

I keep searching the web for a solution and saw that Pytorch Data Loader can make such problems, I tried all the solutions and nothing helped.

Dataset class:


Data Loader:

will be greatly appreciated an answer.