Update modeling_duo_predict_gpt2.py

modeling_duo_predict_gpt2.py

@@ -129,10 +129,10 @@ def sdpa_attention_forward(
         query,
         key,
         value,
-        attn_mask=create_attention_mask_matrix(query.shape[-2]).to(query.device),
+        attn_mask=create_attention_mask_matrix(query.shape[-2]).to(query.device) if module.training else None,
         dropout_p=dropout,
         scale=scaling,
-        is_causal=is_causal,
+        is_causal=False if module.training else True,
     )
     attn_output = attn_output.transpose(1, 2).contiguous()
 
@@ -582,9 +582,12 @@ class DuoPredictGPT2Model(DuoPredictGPT2PretrainedModel):
             inputs_embeds = self.wte(input_ids)
         position_embeds = self.wpe(position_ids)
         ###TODO: correctly initialized
-        hidden_states = torch.empty((batch_size, input_shape[-1], self.embed_dim), device=device)
-        hidden_states[:, ::2] = inputs_embeds[:, ::2] + position_embeds.to(inputs_embeds.device)
-        hidden_states[:, 1::2] = inputs_embeds[:, 1::2] + position_embeds[:, :self.config.max_position_embeddings-1].to(inputs_embeds.device)
+        if inputs_embeds.shape[1] != position_embeds.shape[1]:
+            hidden_states = torch.empty((batch_size, input_shape[-1], self.embed_dim), device=device)
+            hidden_states[:, ::2] = inputs_embeds[:, ::2] + position_embeds.to(inputs_embeds.device)
+            hidden_states[:, 1::2] = inputs_embeds[:, 1::2] + position_embeds[:, :self.config.max_position_embeddings-1].to(inputs_embeds.device)
+        else:
+            hidden_states = inputs_embeds + position_embeds
 
         # Attention mask.
         _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
@@ -897,5 +900,6 @@ if __name__  == "__main__":
     model = DuoPredictGPT2LMHeadModel(cg)
     from src.utils.model_utlis import print_trainable_parameters
     print_trainable_parameters(model)
+    model.eval()
     model(torch.randint(0, 10000, (1, 100)))
     print()