Facing issue with pt_loss computation while evaluation in minillm trainer.

[TanayNarshana]

Thank you very much for sharing the code.

I am trying to run the scripts when I am facing the following issue in minillm/trainer.py script.

Traceback (most recent call last):
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/train_minillm.py", line 99, in <module>
    main()
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/train_minillm.py", line 85, in main
    train(
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/minillm/__init__.py", line 50, in train
    trainer.train()
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/minillm/trainer.py", line 306, in train
    self.evaluate()
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/minillm/trainer.py", line 408, in evaluate
    eval_pt_results = self.evaluate_pt()
  File "/home/ec2-user/SageMaker/tanayn/llm-exploration/knowledge-distillation/minillm/minillm/trainer.py", line 527, in evaluate_pt
    _, stats = self.losses.pt_loss(batch)
TypeError: Loss.pt_loss() missing 1 required positional argument: 'logits'

It seems that the Loss.pt_loss() method requires logits as well. What would be the right way to fix this error?