For now, I think the problem is because I am calling the sub_model2 within a loop.
y = sub_model1(x)
while (not_finished()):
     y1 = pick_some_index(y)
     out = sub_model2(y1)

Based on this https://erickguan.me/2019/pytorch-parallel-model 1, DataParallel copy all model and mini-batch for each device. I guess the copied model in each device might not be deleted (nor handled wisely by Pytorch?) for each iteration in the loop (CMIIW).
I change my training stage by eradicating these loops. It works now with 2 * batch_size.