Returned state if stateful module

neuromorphs · Oct 10, 2023 · 1465f29 · 1465f29
1 parent a0a2f47
commit 1465f29
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
@@ -116,7 +116,10 @@ def forward(self, data: torch.Tensor, state: Optional[Dict[str, Any]] = {}):
                     0
                 )  # Multiple inputs are summed
                 outs[node.name] = self._apply_module(node, input_data, state)
-        return outs[node.name]
+        if len(state) > 0:
+            return outs[node.name], state
+        else:
+            return outs[node.name]
 
 
 def _mod_nir_to_graph(nir_graph: nir.NIRNode) -> Graph:

diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
@@ -79,5 +79,7 @@ def forward(self, x, state=None):
         edges=[("li", "li2")],
     )  # Mock node
     m = load(g, lambda m: StatefulModel())
-    out = m(torch.ones(10))
+    out, state = m(torch.ones(10))
     assert torch.allclose(out, torch.ones(10) * 3)
+    assert state["li"] == (1, )
+    assert state["li"] == (1, )