[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中計算了。其他計算語句不變。