自 PyTorch 1.6 版本开始,F.softmax()
函数不再隐式选择维度,默认情况下不会对最后一个维度进行计算。为了解决这个警告并明确指定要应用 softmax
的维度,你需要修改代码如下:
# 假设 x 是一个形状为 (batch_size, num_classes) 或其他包含多个维度的张量
dim = -1 # 对于最后一维
prediction = F.softmax(x, dim=dim)
这里的 -1
表示最后一个维度,这是通常用于分类任务时 softmax 的应用位置。