使用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’)。