Upload modeling_moment.py

modeling_moment.py

@@ -481,10 +481,12 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         hidden_states = input_mask_patch_view_for_hidden_states * hidden_states
         # [batch_size, n_channels x n_patches, d_model]
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.config.d_model)
-        input_mask_patch_view_for_hidden_states = input_mask_patch_view_for_hidden_states.transpose(1, 2).reshape(batch_size, -1, self.config.d_model)
+        
+        # [batch_size x n_patches]
+        input_mask_patch_view_for_mists = Masking.convert_seq_to_patch_view(input_mask, self.patch_len)
 
         return TimeseriesOutputs(
-            embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states, input_mask_patch_view=input_mask_patch_view_for_hidden_states,
+            embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states, input_mask_patch_view=input_mask_patch_view_for_mists
         )
     
     def forward(