Upload modeling_moment.py

modeling_moment.py

@@ -470,15 +470,14 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         # [batch_size, n_channels x n_patches, d_model]
         # Ensure hidden_states are consistent for both short and long inputs with input_mask specified
         # 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)
-        # [batch_size x n_channels x n_patches x d_model]
-        hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.d_model)
         # [batch_size x n_patches]
         input_mask_patch_view_for_hidden_states = Masking.convert_seq_to_patch_view(input_mask, self.patch_len)
         # [batch_size　x n_channels x n_patches x d_model]
         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]
+        # [batch_size x n_channels x n_patches x d_model]
+        hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.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)