Implement a wrapper for heat flux calculation

[GardevoirX]

The idea is to implement the heat flux calculation as a torch script, and as a result it can be easily used in both ASE and LAMMPS.

The model will be implemented in this way (Thanks @Luthaf ):

class HeatFlux(torch.nn.Module):

    def __init__(self, model: AtomisticModel):
        super().__init__()
        self.model = model.module

    def forward(
        self, 
        systems: List[System], 
        outputs: Dict[str, ModelOutput], 
        selected_atoms = None
    ) -> Dict[str, TensorMap]:
        
        if "experimental::heat_flux" in outputs:
            ..

        else:
            return self.model(systems, outputs, selected_atoms)



wrapper = HeatFlux(pet_model)
wrapper.save("heat_flux.pt")

calc = MetatomicCalculator("heat_flux.pt")
outputs = calc.run_model(atoms, {"experimental::heat_flux": ModelOutput(per_atom=False)})

heat_flux = outputs["experimental::heat_flux"].block().values

And will be located in the models folder under metatomic/python/metatomic_torch/metatomic/torch/.