Upload matryoshka.py

matryoshka.py

@@ -1612,7 +1612,7 @@ class MatryoshkaFusedAttnProcessor1_0_or_2_0:
             hidden_states = attn.group_norm(hidden_states)  # .transpose(1, 2)).transpose(1, 2)
 
         # Reshape hidden_states to 2D tensor
-        hidden_states = hidden_states.view(batch_size, channel, height * width).permute(0, 2, 1).contiguous()
+        hidden_states = hidden_states.view(batch_size, channel, height * width).permute(0, 2, 1)#.contiguous()
         # Now hidden_states.shape is [batch_size, height * width, channels]
 
         if encoder_hidden_states is None:
@@ -1664,8 +1664,8 @@ class MatryoshkaFusedAttnProcessor1_0_or_2_0:
             dropout_p=attn.dropout,
         )
 
-        hidden_states = hidden_states.to(query.dtype)
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
 
         if self_attention_output is not None:
             hidden_states = hidden_states + self_attention_output