Thank you for your reply.
I change the code by following  your advice and get a new error:
/home/wangxu/miniconda3/envs/pt2/lib/python2.7/site-packages/torch/nn/parallel/data_parallel.py:24: UserWarning:
There is an imbalance between your GPUs. You may want to exclude GPU 0 which
has less than 75% of the memory or cores of GPU 1. You can do so by setting
the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES
environment variable.
warnings.warn(imbalance_warn.format(device_ids[min_pos], device_ids[max_pos]))
Traceback (most recent call last):
File “Train_gpu.py”, line 130, in 
train()
File “Train_gpu.py”, line 95, in train
loss1_1 = loss(out1_1 * mask_var, vecmap_var)
File “/home/wangxu/miniconda3/envs/pt2/lib/python2.7/site-packages/torch/nn/modules/module.py”, line 491, in call
result = self.forward(*input, **kwargs)
File “/home/wangxu/miniconda3/envs/pt2/lib/python2.7/site-packages/torch/nn/modules/loss.py”, line 372, in forward
return F.mse_loss(input, target, size_average=self.size_average, reduce=self.reduce)
File “/home/wangxu/miniconda3/envs/pt2/lib/python2.7/site-packages/torch/nn/functional.py”, line 1569, in mse_loss
input, target, size_average, reduce)
File “/home/wangxu/miniconda3/envs/pt2/lib/python2.7/site-packages/torch/nn/functional.py”, line 1537, in _pointwise_loss
return lambd_optimized(input, target, size_average, reduce)
RuntimeError: after cudaLaunch in triple_chevron_launcher::launch(): out of memory