Upload modeling.py

modeling.py

@@ -418,7 +418,7 @@ class MoEGPTForCausalLM(PreTrainedModel):
                 else idx[:, -self.config.sequence_length :]
             )
             # forward the model to get the logits for the index in the sequence
-            logits = self(idx_cond, date, get_logits=True)["logits"]
+            logits = self(idx_cond, date, get_logits=True).logits
             # pluck the logits at the final step and scale by desired temperature
             logits = logits[:, -1, :] / temperature
             # optionally crop the logits to only the top k options