Update CpuLstmModel.py

fixed an interface issue so now it has the same interface as other models.
mhpi · Dec 2, 2023 · 032fbe8 · 032fbe8
1 parent c39a815
commit 032fbe8
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/hydroDL/model/rnn/CpuLstmModel.py b/hydroDL/model/rnn/CpuLstmModel.py
@@ -33,10 +33,9 @@ def forward(self, inputs, doDropMC=False):
         # return out
         results = {}
         x = inputs
-#         x = inputs["xTrain"]
         nt, ngrid, nx = x.shape
         yt = torch.zeros(ngrid, 1)
-        results["yP"] = torch.zeros(nt, ngrid, self.ny)
+        results = torch.zeros(nt, ngrid, self.ny)
         ht = None
         ct = None
         resetMask = True
@@ -46,5 +45,5 @@ def forward(self, inputs, doDropMC=False):
             ht, ct = self.lstm(x0, hidden=(ht, ct), resetMask=resetMask)
             yt = self.linearOut(ht)
             resetMask = False
-            results["yP"][t, :, :] = yt
+            results[t, :, :] = yt
         return results