From d90ea45901e1062d423bf7a9c930bdb0f36f0783 Mon Sep 17 00:00:00 2001 From: Han Yang Date: Thu, 5 Dec 2024 13:42:34 +0800 Subject: [PATCH] update log information --- src/mattersim/forcefield/potential.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/mattersim/forcefield/potential.py b/src/mattersim/forcefield/potential.py index ab87dac..d95fff3 100644 --- a/src/mattersim/forcefield/potential.py +++ b/src/mattersim/forcefield/potential.py @@ -2,7 +2,6 @@ """ Potential """ -import logging import os import pickle import random @@ -19,6 +18,7 @@ from ase.constraints import full_3x3_to_voigt_6_stress from ase.units import GPa from deprecated import deprecated +from loguru import logger from torch.optim import Adam from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR from torch_ema import ExponentialMovingAverage @@ -32,14 +32,6 @@ rank = int(os.getenv("RANK", 0)) -if rank == 0: - logging.basicConfig( - level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" - ) -else: - logging.basicConfig(level=logging.CRITICAL) -logger = logging.getLogger(__name__) - @compile_mode("script") class Potential(nn.Module): @@ -873,6 +865,10 @@ def from_checkpoint( ): load_path = os.path.join(checkpoint_folder, "mattersim-v1.0.0-1M.pth") if not os.path.exists(load_path): + logger.info( + "The pre-trained model is not found locally, " + "attempting to download it from the server." + ) download_checkpoint( "mattersim-v1.0.0-1M.pth", save_folder=checkpoint_folder ) @@ -883,6 +879,10 @@ def from_checkpoint( ): load_path = os.path.join(checkpoint_folder, "mattersim-v1.0.0-5M.pth") if not os.path.exists(load_path): + logger.info( + "The pre-trained model is not found locally, " + "attempting to download it from the server." + ) download_checkpoint( "mattersim-v1.0.0-5M.pth", save_folder=checkpoint_folder )