terminating_states are now calculated using self.all_states

GFNOrg · Nov 15, 2024 · de9edea · de9edea
1 parent 8259a21
commit de9edea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/gfn/gym/hypergrid.py b/src/gfn/gym/hypergrid.py
@@ -178,7 +178,7 @@ def get_terminating_states_indices(self, states: DiscreteStates) -> torch.Tensor
 
         Returns the indices of the terminating states in the canonical ordering as a tensor of shape `batch_shape`.
         """
-        return self.get_states_indices(states)
+        return self.get_states_indices(self.all_states)
 
     @property
     def n_states(self) -> int: