tianle-BigRice:

I don’t understand why the new data set can’t work properly, is it because the input and output are too large? But I also reduced the batchsize and still can’t run


Yes, most likely. If neither a batch size of 1 can run nor you are able to use checkpointing, you could try to use model sharding, i.e. executing separate parts of the model on different GPUs.
If also this needs to much memory, you would have to change the model architecture and make sure the GPU requirement meets the available device memory.