Fix composition tests

alexhernandezgarcia · Oct 24, 2023 · 80a32cf · 80a32cf
1 parent c5106eb
commit 80a32cf
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/gflownet/envs/crystals/composition.py b/gflownet/envs/crystals/composition.py
@@ -425,8 +425,8 @@ def state2oracle(self, state: List = None) -> Tensor:
         if state is None:
             state = self.state
         return self.statetorch2oracle(
-            torch.unsqueeze(tfloat(states, device=self.device), 0)
-        )
+            torch.unsqueeze(tfloat(state, device=self.device, float_type=self.float), 0)
+        )[0]
 
     def statetorch2oracle(
         self, states: TensorType["batch", "state_dim"]

diff --git a/tests/gflownet/envs/test_composition.py b/tests/gflownet/envs/test_composition.py
@@ -38,15 +38,39 @@ def test__environment__initializes_properly(elements):
     [
         (
             [0, 0, 2, 0],
-            [0, 0, 2, 0],
+            [
+                # fmt: off
+                0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                # fmt: on
+            ],
         ),
         (
             [3, 0, 0, 0],
-            [3, 0, 0, 0],
+            [
+                # fmt: off
+                0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                # fmt: on
+            ],
         ),
         (
             [0, 1, 0, 1],
-            [0, 1, 0, 1],
+            [
+                # fmt: off
+                0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                # fmt: on
+            ],
         ),
     ],
 )