black

alexhernandezgarcia · May 7, 2024 · 7e9cc44 · 7e9cc44
1 parent 4267acb
commit 7e9cc44
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 13 deletions.
diff --git a/gflownet/proxy/corners.py b/gflownet/proxy/corners.py
@@ -40,19 +40,16 @@ def min(self):
         return self._min
 
     def __call__(self, states: TensorType["batch", "state_dim"]) -> TensorType["batch"]:
-        return (
-            self.mulnormal_norm
-            * torch.exp(
-                -0.5
-                * (
-                    torch.diag(
+        return self.mulnormal_norm * torch.exp(
+            -0.5
+            * (
+                torch.diag(
+                    torch.tensordot(
                         torch.tensordot(
-                            torch.tensordot(
-                                (torch.abs(states) - self.mu_vec), self.cov_inv, dims=1
-                            ),
-                            (torch.abs(states) - self.mu_vec).T,
-                            dims=1,
-                        )
+                            (torch.abs(states) - self.mu_vec), self.cov_inv, dims=1
+                        ),
+                        (torch.abs(states) - self.mu_vec).T,
+                        dims=1,
                     )
                 )
             )

diff --git a/gflownet/proxy/tetris.py b/gflownet/proxy/tetris.py
@@ -17,7 +17,7 @@ def setup(self, env=None):
     @property
     def norm(self):
         if self.normalize:
-            return (self.height * self.width)
+            return self.height * self.width
         else:
             return 1.0