out of memory(your code)

[LeaRun]

Loading the dataset takes up too much memory, resulting in frequent training stops, 512 × The image size is 512, and the batch size is 8.

[Lei]

You can try load only the dataset list, and load data in each training iteration.