Upload modeling_moment.py

modeling_moment.py

@@ -478,7 +478,10 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         input_mask_patch_view_for_hidden_states = input_mask_patch_view_for_hidden_states.unsqueeze(-1).repeat(
             1, n_channels, 1, self.config.d_model
         )
+        # [batch_size　x n_channels x n_patches x d_model]
         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)
 
         return TimeseriesOutputs(
             embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states