chore: remove jnp.where from _get_ones_like_expanded_block method

instadeepai · Mar 11, 2024 · c08fc62 · c08fc62
1 parent a05b35f
commit c08fc62
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/jumanji/environments/packing/flat_pack/env.py b/jumanji/environments/packing/flat_pack/env.py
@@ -378,15 +378,9 @@ def _is_legal_action(
         return legal
 
     def _get_ones_like_expanded_block(self, grid_block: chex.Array) -> chex.Array:
-        """Makes a grid of zeroes with ones where the block is placed.
+        """Makes a grid of zeroes with ones where the block is placed."""
 
-        Args:
-            grid_with_ones: block placed on a grid of zeroes.
-        """
-
-        grid_with_ones = jnp.where(grid_block != 0, 1, 0)
-
-        return grid_with_ones
+        return (grid_block != 0).astype(jnp.int32)
 
     def _expand_block_to_grid(
         self,