Avoid printing out all the params

bit-bots · Feb 9, 2025 · af70a87 · af70a87
1 parent 048cdca
commit af70a87
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ddlitlab2024/ml/training/distill.py b/ddlitlab2024/ml/training/distill.py
@@ -118,15 +118,15 @@
     teacher_model = End2EndDiffusionTransformer(**model_config).to(device)
 
     # Load the model if a checkpoint is provided
-    logger.info(f"Loading model from {checkpoint}")
+    logger.info(f"Loading teacher model")
     teacher_model.load_state_dict(checkpoint["model_state_dict"])
 
     # Clone the model
     student_model = End2EndDiffusionTransformer(**model_config).to(device)
 
     # Load the same checkpoint into the student model
     # I load it from disk do avoid any potential issues when copying the model
-    logger.info(f"Loading model from {checkpoint}")
+    logger.info(f"Loading student model from teacher checkpoint")
     student_model.load_state_dict(torch.load(args.checkpoint, weights_only=True)["model_state_dict"])
 
     # Create optimizer and learning rate scheduler