RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False

RuntimeError:试图反序列化 CUDA 设备上的对象,但 torch.cuda.is_available() 为 False。如果您 运行 在仅 CPU 的机器上,请使用 torch.load 和 map_location=torch.device('cpu') 将您的存储映射到CPU.

我收到上述代码错误:-

    def get_model(path, device):
        model = models.vgg16(pretrained=False)
    
        for param in model.parameters():
            param.requires_grad = False
        n_inputs = model.classifier[6].in_features
    
        model.classifier[6] = torch.nn.Sequential(
            torch.nn.Linear(n_inputs, 256), torch.nn.ReLU(), torch.nn.Dropout(0.2),
            torch.nn.Linear(256, 10), torch.nn.LogSoftmax(dim=1))
    
        model.load_state_dict(torch.load(path), map_location=torch.device('cpu'))
        model.to(device)
        model.eval()
        return model
    
    device = torch.device("cpu")
    model = get_model('vgg16.pt', device)

您将 map_location 传递给错误的函数(传递给 model.load_state_dict 而不是 torch.load)。

更正后的行如下所示:

model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))