RuntimeError: “nll_loss_forward_reduce_cuda_kernel_2d_index“ not implemented for ‘Int‘解决办法

运行pytorch在计算cross_entropy的loss遇到报错

RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Int'`

原因是pytorch自带的损失计算函数不支持原本的标签1,2,3,4…,n,需要转换成网络需要的one_hot编码才行
具体代码如下

 one_hot = torch.zeros(np.array(batch_size, num_class, device=torch.device('cuda:0')).scatter_(1, label, 1)

batch_size和num_class视自身数据集情况而定,label数据集格式为[batch_size,1],如果不是二维,那么改为label.unsqueeze(1),数据格式要是int64才能被scatter函数调用

你可能感兴趣的:(深度学习,pytorch,深度学习,python)