diff --git a/py/nodes/random_checkpoint_loader.py b/py/nodes/random_checkpoint_loader.py
index 5b80b049..b9c4ff24 100644
--- a/py/nodes/random_checkpoint_loader.py
+++ b/py/nodes/random_checkpoint_loader.py
@@ -171,7 +171,7 @@ class RandomCheckpointLoaderLM:
         ckpt_name: str,
         select_at_random: bool = False,
         base_model: str = "Any",
-    ) -> Tuple[Any, Any, Any]:
+    ) -> Tuple[Any, Any, Any, str]:
         """Load a checkpoint by name, supporting extra folder paths
 
         Args: