feat: print the model's name when logging the number of model paramet…

…ers;
WenjieDu · Dec 20, 2023 · 88056a4 · 88056a4
1 parent 15afa72
commit 88056a4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pypots/base.py b/pypots/base.py
@@ -543,7 +543,8 @@ def _print_model_size(self) -> None:
         """Print the number of trainable parameters in the initialized NN model."""
         num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
         logger.info(
-            f"Model initialized successfully with the number of trainable parameters: {num_params:,}"
+            f"A {self.__class__.__name__} model initialized with the given hyperparameters, "
+            f"the number of trainable parameters: {num_params:,}"
         )
 
     @abstractmethod