Update modules/model.py

modules/model.py

@@ -78,7 +78,7 @@ class CrossAttnProcessor(nn.Module):
         attention_mask=None,
     ):
         batch_size, sequence_length, _ = hidden_states.shape
-        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size=batch_size)
 
         encoder_states = hidden_states
         is_xattn = False