Add custom training loop to the docs

[RaulPPelaez]

Add a more thorough description of the TorchMD_Net model and an example on how to carry out custom training loops.

I hanged the compiled version of this PR here for reviewing purposes https://torchmd-net.readthedocs.io/en/update_docs/models.html

[RaulPPelaez]

@guillemsimeon, @AntonioMirarchi pls review. The code examples run fine, but maybe I missed something.

[guillemsimeon]

Looks good to me. I think in #296 Lewis is interested in just instantiating TorchMD_ET, not the whole TorchMD_Net thing. But for the docs, I find the description very illustrative. Once this is merged, we can refer to this PR and we can tell Lewis that either you design a custom output model to be integrated into the whole TorchMD_Net thing by using create_model, or you can instantiate any representation model, but not using create_model in that case, and taking into account that in that case they get (num_atoms, hidden_channels) scalars and vectors as outputs.