[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

src/neurosis/modules/attention.py

@@ -325,11 +325,13 @@ def forward(
 
         b, _, _ = q.shape
         q, k, v = map(
-            lambda t: t.unsqueeze(3)
-            .reshape(b, t.shape[1], self.heads, self.dim_head)
-            .permute(0, 2, 1, 3)
-            .reshape(b * self.heads, t.shape[1], self.dim_head)
-            .contiguous(),
+            lambda t: (
+                t.unsqueeze(3)
+                .reshape(b, t.shape[1], self.heads, self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(b * self.heads, t.shape[1], self.dim_head)
+                .contiguous()
+            ),
             (q, k, v),
         )
 

src/neurosis/modules/diffusion/model.py

@@ -202,11 +202,13 @@ def attention(self, h_: Tensor) -> Tensor:
         q, k, v = map(lambda x: rearrange(x, "b c h w -> b (h w) c"), (q, k, v))
 
         q, k, v = map(
-            lambda t: t.unsqueeze(3)
-            .reshape(B, t.shape[1], 1, C)
-            .permute(0, 2, 1, 3)
-            .reshape(B * 1, t.shape[1], C)
-            .contiguous(),
+            lambda t: (
+                t.unsqueeze(3)
+                .reshape(B, t.shape[1], 1, C)
+                .permute(0, 2, 1, 3)
+                .reshape(B * 1, t.shape[1], C)
+                .contiguous()
+            ),
             (q, k, v),
         )
         out = xops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)