diff --git a/src/fairchem/core/models/equiformer_v2/equiformer_v2.py b/src/fairchem/core/models/equiformer_v2/equiformer_v2.py index 6171098cb..44cee6430 100644 --- a/src/fairchem/core/models/equiformer_v2/equiformer_v2.py +++ b/src/fairchem/core/models/equiformer_v2/equiformer_v2.py @@ -642,7 +642,7 @@ def forward(self, data: Batch, emb: dict[str, torch.Tensor | GraphData]): return {"energy": energy / data.natoms} else: raise ValueError(f"agg_fn can only be sum or mean, user provided: {self.agg_fn}") - + @registry.register_model("equiformer_v2_force_head")