I would try to check which part of the code is using the majority of the memory and see if this is indeed necessary. E.g. I don’t completely understand your data loading logic, but in case you are trying to preload the samples and store them in images_list, maybe stick to lazily loading the samples.