Don't import torch globally in calculator

nequix/calculator.py

@@ -13,16 +13,14 @@
     dict_to_pytorch_geometric,
     preprocess_graph,
 )
-from nequix.model import Nequix
 from nequix.model import load_model as load_model_jax
 from nequix.model import save_model as save_model_jax
 from nequix.pft.hessian import hessian_linearized
-from nequix.torch_impl.model import NequixTorch
 
 
 def from_pretrained(
     model_name: str = None, model_path: str = None, backend: str = "jax", use_kernel: bool = True
-) -> tuple[Nequix | NequixTorch, dict]:
+):
     """Load a pretrained Nequix model from the cache, checkpoint path, or download it if necessary."""
     assert backend in ["jax", "torch"], f"invalid backend: {backend}"
     base_path = Path("~/.cache/nequix/models/").expanduser()

nequix/pft/hessian.py

@@ -1,6 +1,5 @@
 import jax
 import jax.numpy as jnp
-import equinox as eqx
 
 
 def hessian_linearized(model, graph, batch_size=None):

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nequix"
-version = "0.4.2"
+version = "0.4.3"
 description = "Nequix source code"
 readme = "README.md"
 authors = [