RuntimeError: Expected a ‘cuda‘ device type for generator but found ‘cpu‘

        使用pytoch1.10训练时,报错:RuntimeError: Expected a ‘cuda‘ device type for generator but found ‘cpu‘。同样的代码,在当前机器上训练可以,再换到另一个机器上报错。网上查解决方式有几种,我是通过第一种方式解决的,这里整理下:

1、修改DataLoader的generator(问题解决)

将:

        data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True)

改为:

        data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True,generator=torch.Generator(device = 'cuda'))

2、修改将torch.set_default_tensor_type

        将torch.set_default_tensor_type(‘torch.FloatTensor’)改为:torch.set_default_tensor_type(‘torch.cuda.FloatTensor’)。

PyTorch 有哪些坑/bug? - 知乎

你可能感兴趣的:(未分类,java,数据库,前端)