docstring

ACEsuit · May 12, 2023 · 5edaff9 · 5edaff9
1 parent 59b30f0
commit 5edaff9
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/mace_jax/tools/train.py b/mace_jax/tools/train.py
@@ -100,6 +100,14 @@ def evaluate(
     data_loader: data.GraphDataLoader,
     name: str = "Evaluation",
 ) -> Tuple[float, Dict[str, Any]]:
+    r"""Evaluate the model on the given data loader.
+
+    Args:
+        model: function of signature `model(params, graph) -> {energy: [num_graphs], forces: [num_nodes, 3], stress: [num_graphs, 3, 3]}`
+        params: parameters of the model
+        loss_fn: function of signature `loss_fn(graph, output) -> loss` where `output` is the output of `model`
+        data_loader: data loader
+    """
     total_loss = 0.0
     num_graphs = 0