diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py
index 33621bc3..60916a4e 100644
--- a/src/aare/daq/mlbox.py
+++ b/src/aare/daq/mlbox.py
@@ -359,7 +359,7 @@ class MlBox:
                 logger.debug(f"Failed to parse prediction: {e}", exc_info=True)
                 continue
 
-            base_key = MLOutputModel.get_class_str(cls=cls)
+            base_key = MLOutputModel.get_class_str(klass=cls)
             if base_key not in out.boxes:
                 out.boxes[base_key] = MLBoxModel(
                     cls=cls,