Upload modeling_moment.py

modeling_moment.py

@@ -449,12 +449,15 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         outputs = self.encoder(inputs_embeds=enc_in, attention_mask=attention_mask)
         enc_out = outputs.last_hidden_state
 
+        # For Mists model
+        hidden_states = outputs.last_hidden_state
+
         enc_out = enc_out.reshape((-1, n_channels, n_patches, self.config.d_model))
         # [batch_size x n_channels x n_patches x d_model]
 
         # 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 = enc_out.reshape(batch_size, n_channels * n_patches, self.config.d_model)
 
         if reduction == "mean":
             enc_out = enc_out.mean(dim=1, keepdim=False)  # Mean across channels