Update crit.py

mhpi · Dec 1, 2023 · 8d87998 · 8d87998
1 parent f701b32
commit 8d87998
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/hydroDL/model/crit.py b/hydroDL/model/crit.py
@@ -202,8 +202,7 @@ def forward(self, output, target, igrid):
         nt = target.shape[0]
         stdse = np.tile(self.std[igrid].T, (nt, 1))
 
-        tdbatch = torch.tensor(stdse, requires_grad=False).float().to(device)
-
+        stdbatch = torch.tensor(stdse, requires_grad=False).float().to(self.device)
         p0 = output[:, :, 0]   # dim: Time*Gage
         t0 = target[:, :, 0]
         mask = t0 == t0
@@ -232,7 +231,7 @@ def __init__(self, stdarray, eps=0.1,device = torch.cuda.current_device()):
     def forward(self, output, target, igrid):
         nt = target.shape[0]
         stdse = np.tile(self.std[igrid], (nt, 1))
-        tdbatch = torch.tensor(stdse, requires_grad=False).float().to(device)
+        stdbatch = torch.tensor(stdse, requires_grad=False).float().to(self.device)
         p0 = output[:, :, 0]   # dim: Time*Gage
         t0 = target[:, :, 0]
         mask = t0 == t0