报错:RuntimeError: Expected object of backend CUDA but got backend CPU for argument #3 'index'

日萌社

人工智能AI:Keras PyTorch MXNet TensorFlow PaddlePaddle 深度学习实战(不定时更新)


    RuntimeError: Expected object of backend CUDA but got backend CPU for argument #3 'index'
    分析:data数据或者model没有调用.to(device)
    解决:
        # 进行可用设备检测, 有GPU的话将优先使用GPU
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device) 或 model = model.cuda(device)
        output = model(input.to(device)) 或 output = model(input.cuda(device))
        # 选择损失函数, 这里选择预定义的交叉熵损失函数
        criterion = torch.nn.CrossEntropyLoss().to(device) 或 criterion = torch.nn.CrossEntropyLoss().cuda(device)
        loss = criterion(output, lable.to(device)) 或 loss = criterion(output, lable.cuda(device))
发布了452 篇原创文章 · 获赞 154 · 访问量 23万+

猜你喜欢

转载自blog.csdn.net/zimiao552147572/article/details/105496985
今日推荐