Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you

一、问题

使用了自己创建的dataset和collate_fn，使用了默认的dataloader，当设置num_worker的时候，超过0则报如下错误：

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

二、原因

查看了平时其他人的代码，发现自己在collate_fn对tensor使用了.to(device)，而别人则是在从dataloader取出数据后，放入模型前，将其移到device中，例如：

    for epoch in range(max_epoch):
        for batch_idx, (en_id, att_mask_en, zh_id, att_mask_zh) in enumerate(data_loader): 
                en_id = en_id.to(device)

而我则是错误地在collate_fn中将数据移至device中：

def collate_fn(batch):
    en, zh = list(zip(*batch))
    en_output = tokenizer.encode_batch(en,is_pretokenized=True)
    zh_id,att_mask_zh = zh_mapper.encode_batch(zh)
    en_id =[]
    att_mask_en = []
    for item in en_output:
        en_id.append(item.ids)
        att_mask_en.append(item.attention_mask)
    en_id = torch.tensor(en_id, dtype=torch.long).to(device)
    att_mask_en = torch.tensor(att_mask_en, dtype=torch.bool).to(device)
    zh_id = torch.tensor(zh_id, dtype=torch.long).to(device)
    att_mask_zh = torch.tensor(att_mask_zh, dtype=torch.bool).to(device)

    return en_id, att_mask_en, zh_id, att_mask_zh

当然，在定义dataset，方法__getitem__ 里把张量放到device里也会产生同样的错误。

此外，我还使用了官方建议的spawn:

if __name__ == '__main__':
    torch.multiprocessing.set_start_method('spawn')
    train()

但是gpu会卡住不运行数据加载。