I simplify my project to a file segment.py:
#!/usr/bin/env python

import torch
from torch import nn
import torch.utils.data


class DRN(nn.Module):
    def __init__(self):
        super(DRN, self).__init__()
        self.a = nn.Conv2d(3, 16, kernel_size=7)

    def forward(self, x):
        print('before DRN forward')
        return x


if __name__ == '__main__':
    model = torch.nn.DataParallel(DRN()).cuda().train()
    input_ = torch.rand(2).cuda()
    print('before input')
    model(input_)
    print('end')

If I run CUDA_LAUNCH_BLOCKING=1 CUDA_VISIBLE_DEVICES=0,1 ./segment.py, It will stucks after print before input.
However, if I change rand(2) to rand(1), it does not stuck again.