changed default trainer to energy (#298)

Co-authored-by: Muhammed Shuaibi <45150244+mshuaibii@users.noreply.github.com>
FAIR-Chem · Oct 18, 2021 · 458da25 · 458da25
1 parent 1044e31
commit 458da25
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 7 deletions.
diff --git a/main.py b/main.py
@@ -40,7 +40,7 @@ def __call__(self, config):
         try:
             setup_imports()
             self.trainer = registry.get_trainer_class(
-                config.get("trainer", "simple")
+                config.get("trainer", "energy")
             )(
                 task=config["task"],
                 model=config["model"],

diff --git a/ocpmodels/common/relaxation/ase_utils.py b/ocpmodels/common/relaxation/ase_utils.py
@@ -95,7 +95,7 @@ def __init__(
         self.config["checkpoint"] = checkpoint
 
         self.trainer = registry.get_trainer_class(
-            config.get("trainer", "simple")
+            config.get("trainer", "energy")
         )(
             task=config["task"],
             model=config["model"],

diff --git a/ocpmodels/models/gemnet/fit_scaling.py b/ocpmodels/models/gemnet/fit_scaling.py
@@ -20,11 +20,11 @@
 import torch
 from tqdm import trange
 
-from ocpmodels.models.gemnet.layers.scaling import AutomaticFit
-from ocpmodels.models.gemnet.utils import write_json
 from ocpmodels.common.flags import flags
 from ocpmodels.common.registry import registry
 from ocpmodels.common.utils import build_config, setup_imports, setup_logging
+from ocpmodels.models.gemnet.layers.scaling import AutomaticFit
+from ocpmodels.models.gemnet.utils import write_json
 
 if __name__ == "__main__":
     setup_logging()
@@ -74,7 +74,7 @@ def initialize_scale_file(scale_file):
 
     AutomaticFit.set2fitmode()
 
-    trainer = registry.get_trainer_class(config.get("trainer", "simple"))(
+    trainer = registry.get_trainer_class(config.get("trainer", "energy"))(
         task=config["task"],
         model=config["model"],
         dataset=config["dataset"],

diff --git a/scripts/hpo/run_tune.py b/scripts/hpo/run_tune.py
@@ -14,7 +14,7 @@
 def ocp_trainable(config, checkpoint_dir=None):
     setup_imports()
     # trainer defaults are changed to run HPO
-    trainer = registry.get_trainer_class(config.get("trainer", "simple"))(
+    trainer = registry.get_trainer_class(config.get("trainer", "energy"))(
         task=config["task"],
         model=config["model"],
         dataset=config["dataset"],

diff --git a/scripts/hpo/run_tune_pbt.py b/scripts/hpo/run_tune_pbt.py
@@ -17,7 +17,7 @@ def ocp_trainable(config, checkpoint_dir=None):
     # update config for PBT learning rate
     config["optim"].update(lr_initial=config["lr"])
     # trainer defaults are changed to run HPO
-    trainer = registry.get_trainer_class(config.get("trainer", "simple"))(
+    trainer = registry.get_trainer_class(config.get("trainer", "energy"))(
         task=config["task"],
         model=config["model"],
         dataset=config["dataset"],