Upload matryoshka.py

matryoshka.py

@@ -1519,7 +1519,8 @@ class MatryoshkaTransformerBlock(nn.Module):
 
         # attn_output_cond = attn_output_cond.permute(0, 2, 1).contiguous()
         attn_output_cond = self.proj_out(attn_output_cond)
-        attn_output_cond = attn_output_cond.permute(0, 2, 1).reshape(batch_size, channels, *spatial_dims)
+        # attn_output_cond = attn_output_cond.permute(0, 2, 1).reshape(batch_size, channels, *spatial_dims)
+        attn_output_cond = attn_output_cond.reshape(batch_size, channels, *spatial_dims)
         hidden_states = hidden_states + attn_output_cond
 
         if self.ff is not None: