Make GraphNeT models "jittable"

Created by: asogaard

Making models "jittable" makes it possible to save model artifacts completely indepent of framework code, in contrast to save the pytorch model and/or state_dict. This requires refactoring all Model classes to make them compatible with torch.jit.script(model).

See:

• https://github.com/pyg-team/pytorch_geometric/pull/1191
• https://github.com/pyg-team/pytorch_geometric/pulls?q=is%3Apr+author%3Algray