Correctly handle cell state

DeepTrackAI · May 23, 2024 · f37b717 · f37b717
1 parent ec8a77d
commit f37b717
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deeplay/models/recurrentmodel.py b/deeplay/models/recurrentmodel.py
@@ -167,7 +167,10 @@ def forward(self, x):
             x = self.embedding_dropout(x)
 
         outputs = x
-        outputs, hidden = super().forward(outputs)
+        if self.return_cell_state:
+            outputs, hidden = super().forward(outputs)
+        else:
+            outputs = super().forward(outputs)
 
         if self.bidirectional:
             outputs = (