Use _get_state(state) and _get_done(done)

alexhernandezgarcia · Nov 23, 2023 · 3961398 · 3961398
1 parent fcfa070
commit 3961398
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/gflownet/envs/crystals/miller.py b/gflownet/envs/crystals/miller.py
@@ -51,10 +51,8 @@ def get_mask_invalid_actions_forward(
             - True if the forward action is invalid from the current state.
             - False otherwise.
         """
-        if state is None:
-            state = self.state.copy()
-        if done is None:
-            done = self.done
+        state = self._get_state(state)
+        done = self._get_done(done)
 
         mask = super().get_mask_invalid_actions_forward(state, done)