多GPU Pytork预测

Question

我试图使用多个GPU来处理一个GPU已经正常运行的代码。 我做了这些修改，使其能够对多个GPU进行di：

classifier = AnnotationCNN.load_from_checkpoint(classifier_ckpt,normalizer_ckpt=normalizer_ckpt)
classifier=classifier.eval()
    
**device = torch.device("cuda" if torch.cuda.is_available() else "cpu")**
classifier = classifier.to(torch.float16).to(device)
**classifier= nn.DataParallel(classifier)**


dataset = InitImageTifDataset(
                data_root=data_root,
                level=level,
                region_width=region_width,
                region_height=region_height,
                crop_size=crop_size,
                image_dir=image_dir,
                image_name=image_name,
                channels=channels
            )

        dl = DataLoader(dataset, batch_size=5, shuffle=False, num_workers=5)
        for i,(img,mean_stddev,filename) in enumerate(dl):
            mask = mean_stddev > 20

            img=img[torch.nonzero(mask)]
            if img.nelement() != 0 :
                **tensor=img[:,0,0,:,:,:].to(device).to(torch.float16).to(device)**
                *_,preds = classifier(tensor)*

但我有以下错误：

RuntimeError: Expected tensor for argument #1 'output' to have the same device as tensor for argument #5 'weight'; but device 0 does not equal 1 (while checking arguments for batch_norm_cuda)

我很确定错误来自数据加载器，它没有适应多GPU。我怎样才能修好它，谢谢你