I try to find whether this error is cause by my model.py, so I changed my model to a very simple CNN network.
Here is my code:
import torch
import torch.nn as nn

class DemoModel(nn.Module):
    def __init__(self):
        super(DemoModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 7, 2, padding=1)
        self.conv2 = nn.Conv2d(64, 64, 3, 2, padding=1)
        self.conv3 = nn.Conv2d(64, 3, 1, padding=1)

    def forward(self, x, not_used1, not_used2):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        return x, out

model = DemoModel()
model = model.to(device)
model = nn.DataParallel(model, device_ids=range(opt.n_GPUs))

However, it will also cause error using two GPUs.
RuntimeError: The size of tensor a (8) must match the size of tensor b (16) at non-singleton dimension 0

Will this error be related to h5 files? Because I read my data from a .h5 file.