def METHOD_NAME(self, trainer, pl_module):
    try:
        npars = pl_module.model.model.num_parameters()
    except AttributeError:
        npars = pl_module.model.num_parameters()
    n_trainable_pars = count_trainable_parameters(pl_module)
    # mp stands for million parameters
    trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6})