Use MetatomicCalculator as the class name in torch profiler

Author: Luthaf

python/examples/4-profiling.py

@@ -183,15 +183,15 @@ def forward(
 # with the same function name as the corresponding torch functions (e.g.
 # ``aten::arange`` is :py:func:`torch.arange`). We can also see some internal functions
 # from metatomic, with the name staring with ``AtomisticModel::`` for
-# :py:class:`AtomisticModel`; and ``ASECalculator::`` for
+# :py:class:`AtomisticModel`; and ``MetatomicCalculator::`` for
 # :py:class:`ase_calculator.MetatomicCalculator`.
 #
 # If you want to see more details on the internal steps taken by your model, you can add
 # :py:func:`torch.profiler.record_function`
 # (https://pytorch.org/docs/stable/generated/torch.autograd.profiler.record_function.html)
 # inside your model code to give names to different steps in the calculation. This is
 # how we are internally adding names such as ``Model::forward`` or
-# ``ASECalculator::prepare_inputs`` above.
+# ``MetatomicCalculator::prepare_inputs`` above.
 #
 
 # %%

python/metatomic_torch/metatomic/torch/ase_calculator.py

@@ -327,7 +327,7 @@ def calculate(
         if "stresses" in properties:
             raise NotImplementedError("'stresses' are not implemented yet")
 
-        with record_function("ASECalculator::prepare_inputs"):
+        with record_function("MetatomicCalculator::prepare_inputs"):
             outputs = self._ase_properties_to_metatensor_outputs(
                 properties,
                 calculate_forces=calculate_forces,
@@ -371,7 +371,7 @@ def calculate(
                 selected_atoms=None,
             )
 
-        with record_function("ASECalculator::compute_neighbors"):
+        with record_function("MetatomicCalculator::compute_neighbors"):
             # convert from ase.Atoms to metatomic.torch.System
             system = System(types, positions, cell, pbc)
 
@@ -394,7 +394,7 @@ def calculate(
         )
         energy = outputs["energy"]
 
-        with record_function("ASECalculator::sum_energies"):
+        with record_function("MetatomicCalculator::sum_energies"):
             if run_options.outputs["energy"].per_atom:
                 assert len(energy) == 1
                 assert energy.sample_names == ["system", "atom"]
@@ -409,11 +409,11 @@ def calculate(
             assert len(energy.block().gradients_list()) == 0
             assert energy.block().values.shape == (1, 1)
 
-        with record_function("ASECalculator::run_backward"):
+        with record_function("MetatomicCalculator::run_backward"):
             if do_backward:
                 energy.block().values.backward()
 
-        with record_function("ASECalculator::convert_outputs"):
+        with record_function("MetatomicCalculator::convert_outputs"):
             self.results = {}
 
             if calculate_energies: