logging

Author: epwalsh

src/olmo_core/distributed/checkpoint/__init__.py

@@ -201,6 +201,9 @@ def load_model_and_optim_state(
         metadata = reader.read_metadata()
         for current_key, original_key in key_mapping.items():
             if original_key in metadata.state_dict_metadata:
+                log.info(
+                    f"Mapping current key '{current_key}' to key '{original_key}' in checkpoint"
+                )
                 state_dict[original_key] = state_dict.pop(current_key)
 
     dist_cp.load(