Merge pull request #165 from GFNOrg/step_type_checking

checks whether user-defined env.step method returns the expected type
GFNOrg · Feb 25, 2024 · 7996b37 · 7996b37
2 parents ccfa959 + 1a5ad2c
commit 7996b37
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/gfn/env.py b/src/gfn/env.py
@@ -219,6 +219,10 @@ def _step(
         not_done_actions = actions[~new_sink_states_idx]
 
         new_not_done_states_tensor = self.step(not_done_states, not_done_actions)
+        if not isinstance(new_not_done_states_tensor, torch.Tensor):
+            raise Exception(
+                "User implemented env.step function *must* return a torch.Tensor!"
+            )
 
         new_states.tensor[~new_sink_states_idx] = new_not_done_states_tensor