System Info

• Transformer version: 4.46.0
• Model: nvidia/Mistral-NeMo-Minitron-8B-Base

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

When calling the forward method on the NeMo Mistral model, the following exception occurs:

[rank2]:   File "/lustre/fsw/portfolios/llmservice/users/gheinrich/anaconda3/envs/vila/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1582, in _call_impl
[rank2]:     result = forward_call(*args, **kwargs)
[rank2]: TypeError: MistralForCausalLM.forward() got an unexpected keyword argument 'num_items_in_batch'

Expected behavior

The forward() method should use num_items_in_batch for the loss calculation.