Update crit.py for CPU LSTM

mhpi · Dec 1, 2023 · bda8a6e · bda8a6e
1 parent b03a83b
commit bda8a6e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/hydroDL/model/crit.py b/hydroDL/model/crit.py
@@ -200,7 +200,10 @@ def __init__(self, stdarray, eps=0.1):
     def forward(self, output, target, igrid):
         nt = target.shape[0]
         stdse = np.tile(self.std[igrid].T, (nt, 1))
-        stdbatch = torch.tensor(stdse, requires_grad=False).float().cuda()
+        if torch.cuda.is_available():
+            stdbatch = torch.tensor(stdse, requires_grad=False).float().cuda()
+        else:
+            stdbatch = torch.tensor(stdse, requires_grad=False).float()
         p0 = output[:, :, 0]   # dim: Time*Gage
         t0 = target[:, :, 0]
         mask = t0 == t0
@@ -227,7 +230,10 @@ def __init__(self, stdarray, eps=0.1):
     def forward(self, output, target, igrid):
         nt = target.shape[0]
         stdse = np.tile(self.std[igrid], (nt, 1))
-        stdbatch = torch.tensor(stdse, requires_grad=False).float().cuda()
+        if torch.cuda.is_available():
+            stdbatch = torch.tensor(stdse, requires_grad=False).float().cuda()
+        else:
+            stdbatch = torch.tensor(stdse, requires_grad=False).float()
         p0 = output[:, :, 0]   # dim: Time*Gage
         t0 = target[:, :, 0]
         mask = t0 == t0