l;l

modelLM.py

@@ -33,13 +33,13 @@ class OBILanguageModel(PreTrainedModel):
         tok_emb = self.token_embedding_table(idx)
         pos_emb = None  # Initialize pos_emb to None
         try:
-            pos_emb = self.position_embedding_table(torch.arange(5, device='cpu'))
+            pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
         except IndexError as e:
             # Handle the IndexError by initializing pos_emb with zeros
             print(f"IndexError: {e}")
-            print(f"idx.size(0): {idx.size(0)}")
+            print(f"idx.size(1): {idx.size(1)}")
             print(f"Positional embedding table shape: {self.position_embedding_table.weight.shape}")
-            pos_emb = torch.zeros((5, self.config.hidden_size), device=device)
+            pos_emb = torch.zeros((idx.size(1), self.config.hidden_size), device=device)
 
         x = tok_emb + pos_emb
         x = self.transformer(x, x)