Upload modeling_moment.py

modeling_moment.py

@@ -456,6 +456,7 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         # For Mists model
         # [batch_size, n_channels x n_patches, d_model]
         # hidden_states = enc_out.reshape(batch_size, n_channels * n_patches, self.config.d_model)
+        hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.d_model).transpose(1, 2).reshape(batch_size, -1, self.config.d_model)
 
         if reduction == "mean":
             enc_out = enc_out.mean(dim=1, keepdim=False)  # Mean across channels