Update crit.py

mhpi · Dec 1, 2023 · c39a815 · c39a815
1 parent 8d87998
commit c39a815
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/hydroDL/model/crit.py b/hydroDL/model/crit.py
@@ -192,7 +192,7 @@ def forward(self, output, target):
 class NSELossBatch(torch.nn.Module):
     # Same as Fredrick 2019, batch NSE loss
     # stdarray: the standard deviation of the runoff for all basins
-    def __init__(self, stdarray,  eps=0.1,device = torch.cuda.current_device()):
+    def __init__(self, stdarray,  eps=0.1,device = 'cpu'):
         super(NSELossBatch, self).__init__()
         self.std = stdarray
         self.eps = eps
@@ -222,7 +222,7 @@ def forward(self, output, target, igrid):
 class NSESqrtLossBatch(torch.nn.Module):
     # Same as Fredrick 2019, batch NSE loss, use RMSE and STD instead
     # stdarray: the standard deviation of the runoff for all basins
-    def __init__(self, stdarray, eps=0.1,device = torch.cuda.current_device()):
+    def __init__(self, stdarray, eps=0.1,device = 'cpu'):
         super(NSESqrtLossBatch, self).__init__()
         self.std = stdarray
         self.eps = eps