[pytorch] 修改cpu计算型loss到gpu上计算

需要改进的代码

tp = (y_true * y_pred).sum(dim=0).to(torch.float32)
tn = ((1 - y_true) * (1 - y_pred)).sum(dim=0).to(torch.float32)
fp = ((1 - y_true) * y_pred).sum(dim=0).to(torch.float32)
fn = (y_true * (1 - y_pred)).sum(dim=0).to(torch.float32)

正常运行，就会报错

(Pdb) torch.sum((y_true * y_pred), dim=0).to(torch.float32)
*** RuntimeError: expected backend CPU and dtype Float but got backend CUDA and dtype Float

我把.sum还成了torch.sum()之后错误没变。
其实做法非常简单：

y_true = make_one_hot(y_true, 2).to(torch.float32).cuda()
y_pred = F.softmax(y_pred, dim=1).cuda()

把数据只要提前送进cuda中，后续的计算就是在gpu中计算了。其他计算语句不变。