RuntimeError: CUDA error (10): invalid device ordinal
两张卡跑的train,在本机一张卡上test不了
结果很简单,直接将 model.module
拿出来就行了
因为之前双卡用的parallel
def __init__(self, module, device_ids=None, output_device=None, dim=0):
super(DataParallel, self).__init__()
if not torch.cuda.is_available():
self.module = module
self.device_ids = []
return
if device_ids is None:
device_ids = list(range(torch.cuda.device_count()))
if output_device is None:
output_device = device_ids[0]
详见:parallel源码
https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html#DataParallel
还没有评论,来说两句吧...