From 80d5d9299b233bb388af4e61ca35f7b2ab5297ce Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 8 Dec 2023 17:23:51 +0100
Subject: [PATCH] Temp fix for mask shape

---
 PixArt/models/PixArt_blocks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/PixArt/models/PixArt_blocks.py b/PixArt/models/PixArt_blocks.py
index 05a9a87..9e492d6 100644
--- a/PixArt/models/PixArt_blocks.py
+++ b/PixArt/models/PixArt_blocks.py
@@ -73,7 +73,7 @@ class MultiHeadCrossAttention(nn.Module):
             q, k, v = map(lambda t: t.permute(0, 2, 1, 3),(q, k, v),)
             
             attn_mask = None
-            if mask is not None:
+            if mask is not None and len(mask) > 1:
                 # This is probably wrong
                 attn_mask = torch.zeros(
                     [1, q.shape[1], q.shape[2], v.shape[2]],