Upload modeling_moment.py

modeling_moment.py

@@ -430,6 +430,8 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         if input_mask is None:
             input_mask = torch.ones((batch_size, seq_len)).to(x_enc.device)
 
+        print("＊input_mask: ", input_mask.shape)
+
         x_enc = self.normalizer(x=x_enc, mask=input_mask, mode="norm")
         x_enc = torch.nan_to_num(x_enc, nan=0, posinf=0, neginf=0)
 
@@ -472,6 +474,7 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         # 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_patches]
+        print("＊input_mask: ", input_mask.shape)
         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(
@@ -480,7 +483,7 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         # [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)
         print("＊input_mask_patch_view_for_hidden_states: ", input_mask_patch_view_for_hidden_states.shape)
-        print("hidden_states: ", hidden_states.shape)
+        print("＊hidden_states: ", hidden_states.shape)
         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)