checks whether user-defined function returns the expected type

GFNOrg · Feb 24, 2024 · 1a5ad2c · 1a5ad2c
1 parent ccfa959
commit 1a5ad2c
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