docstring fix

maanug-nv · maanug-nv · commit b31d7f91f7fd · 2025-03-19T19:59:48.000-07:00
Signed-off-by: Maanu Grover &lt;maanug@nvidia.com&gt;
diff --git a/nemo/tron/llm/gpt.py b/nemo/tron/llm/gpt.py
@@ -44,6 +44,7 @@ def forward_step(state: GlobalState, data_iterator: Iterable, model: GPTModel):
     """Forward training step.
 
     Args:
+        state (GlobalState): Global state for the run
         data_iterator : Input data iterator
         model (GPTModel): The GPT Model
     """
