trainingblock pass param

microsoft · Feb 28, 2024 · fa0ad16 · fa0ad16
1 parent da7be5c
commit fa0ad16
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/orttraining/orttraining/python/training/artifacts.py b/orttraining/orttraining/python/training/artifacts.py
@@ -141,7 +141,7 @@ def build(self, *inputs_to_loss):
 
             return self._loss(*inputs_to_loss)
 
-    training_block = _TrainingBlock(loss_block)
+    training_block = _TrainingBlock(loss_block, loss_input_names)
 
     if requires_grad is not None and frozen_params is not None and set(requires_grad).intersection(set(frozen_params)):
         raise RuntimeError(