Is there any reason why you load all of your images during init ?
Depending on the number of images, it can explain why the execution requires so much memory. You should rather provide an index based on the image address, and load the image while calling _getitem_. The purpose of using a Dataset is to load the data on demand, thus not requiring all the images to be loaded in memory (that’s what you are doing).
Also you can make sure only one copy of your tensor is shared across workers by using .share_memory_() on the underlying tensors, if you have. For example, if you had created a tensor of positions during init, then you would want that tensor to be shared while the Dataloader's workers create the batches.