removed index from BharatAIs forward method and replaced index everywhere with input_ids

model.py

@@ -116,12 +116,13 @@ class BharatAI(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
-    def forward(self, input_ids, index, labels=None): #, targets
-        B, T = index.shape
+    def forward(self, input_ids,labels=None): #, targets # index, 
+        #B, T = index.shape
+        B, T = input_ids.shape
         x = input_ids
 
         # idx and targets are both (B,T) tensor of integers
-        tok_emb = self.token_embedding_table(index) # (B,T,C)
+        tok_emb = self.token_embedding_table(input_ids) # (B,T,C)
         pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)
         x = tok_emb + pos_emb # (B,T,C)
         x = self.blocks(x) # (B,T,C)