Upload modeling_moment.py

modeling_moment.py

@@ -481,9 +481,10 @@ 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)
 
         return TimeseriesOutputs(
-            embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states
+            embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states, input_mask_patch_view=input_mask_patch_view_for_hidden_states,
         )
     
     def forward(