minor fixes (#615)

FAIR-Chem · Jan 9, 2024 · f83c772 · f83c772
1 parent 7732461
commit f83c772
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 329 deletions.
diff --git a/configs/is2re/example.yml b/configs/is2re/example.yml
diff --git a/configs/s2ef/example.yml b/configs/s2ef/example.yml
diff --git a/ocpmodels/trainers/base_trainer.py b/ocpmodels/trainers/base_trainer.py
@@ -437,7 +437,6 @@ def load_model(self) -> None:
         if self.logger is not None:
             self.logger.watch(self.model)
 
-        self.model.to(self.device)
         if distutils.initialized() and not self.config["noddp"]:
             self.model = DistributedDataParallel(
                 self.model, device_ids=[self.device]

diff --git a/ocpmodels/trainers/ocp_trainer.py b/ocpmodels/trainers/ocp_trainer.py
@@ -41,8 +41,11 @@ class OCPTrainer(BaseTrainer):
     Args:
         task (dict): Task configuration.
         model (dict): Model configuration.
+        outputs (dict): Output property configuration.
         dataset (dict): Dataset configuration. The dataset needs to be a SinglePointLMDB dataset.
         optimizer (dict): Optimizer configuration.
+        loss_fns (dict): Loss function configuration.
+        eval_metrics (dict): Evaluation metrics configuration.
         identifier (str): Experiment identifier that is appended to log directory.
         run_dir (str, optional): Path to the run directory where logs are to be saved.
             (default: :obj:`None`)
@@ -60,6 +63,7 @@ class OCPTrainer(BaseTrainer):
             (default: :obj:`False`)
         slurm (dict): Slurm configuration. Currently just for keeping track.
             (default: :obj:`{}`)
+        noddp (bool, optional): Run model without DDP.
     """
 
     def __init__(