revert: back from multiplicative masking to -inf

This reverts commit b3ee035.

src/gflownet/envs/graph_building_env.py

@@ -613,12 +613,8 @@ def _apply_action_masks(self):
         )
 
     def _mask(self, x, m):
-        """
-        mask logit vector x with binary mask m, -1000 is a tiny log-value
-        Note to self: we can't use torch.inf here, because inf * 0 is nan
-        """
         assert m.dtype == torch.float
-        return x * m + -1000 * (1 - m)
+        return x.masked_fill(m == 0., -torch.inf)
 
     def detach(self):
         new = copy.copy(self)
@@ -756,25 +752,8 @@ def sample(self) -> List[ActionIndex]:
         u = [torch.rand(i.shape, device=self.dev) for i in self._masked_logits]
         # Gumbel noise
         gumbel = [logit - (-noise.log()).log() for logit, noise in zip(self._masked_logits, u)]
-
-        if self._action_masks is not None:
-            gumbel_safe = [
-                torch.where(
-                    mask == 1,
-                    torch.maximum(
-                        x,
-                        torch.nextafter(
-                            torch.tensor(torch.finfo(x.dtype).min, dtype=x.dtype), torch.tensor(0.0, dtype=x.dtype)
-                        ).to(x.device),
-                    ),
-                    torch.finfo(x.dtype).min,
-                )
-                for x, mask in zip(gumbel, self._action_masks)
-            ]
-        else:
-            gumbel_safe = gumbel
         # Take the argmax
-        return self.argmax(x=gumbel_safe)
+        return self.argmax(x=gumbel)
 
     def argmax(
         self,