When I run the train.py file, setting the batchsize to 2 will increase the running memory by more than two times. What's going on?