feat: keep num_params as an attribute when run _print_model_size();

pypots/base.py

@@ -497,17 +497,20 @@ def __init__(
  self.num_workers = num_workers
 
  self.model = None
+ self.num_params = None
  self.optimizer = None
  self.best_model_dict = None
  self.best_loss = float("inf")
  self.best_epoch = -1
 
  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)
+ self.num_params = sum(
+ p.numel() for p in self.model.parameters() if p.requires_grad
+ )
  logger.info(
  f"{self.__class__.__name__} initialized with the given hyperparameters, "
- f"the number of trainable parameters: {num_params:,}"
+ f"the number of trainable parameters: {self.num_params:,}"
  )
 
  @abstractmethod