And I found that when batch size=2, I increased the training data from 200 to 550. When I turned on the train loop, I found that the first time from 1686 to 4848, print(batch_idx) will output 0 ,1,2,274, and then the GPU will become 4848MB. When batch size=4, the output will be 0,1,2,3,137.
I think that GPU memory will increase only when all the data is running.
image521×527 80.4 KB