From 02a68ff71ddf051deaf03c7ad604d28b25e2bde9 Mon Sep 17 00:00:00 2001
From: Teddy Koker <teddy.koker@gmail.com>
Date: Fri, 6 Mar 2026 09:20:59 -0500
Subject: [PATCH 1/2] remove torch import from top of file

---
 nequix/calculator.py  | 4 +---
 nequix/pft/hessian.py | 1 -
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/nequix/calculator.py b/nequix/calculator.py
index 302648f..c20e29d 100644
--- a/nequix/calculator.py
+++ b/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()
diff --git a/nequix/pft/hessian.py b/nequix/pft/hessian.py
index c61d3f1..92ad83e 100644
--- a/nequix/pft/hessian.py
+++ b/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):

From 6bb67462d41c11dc80b2e8d89202cb84175096c4 Mon Sep 17 00:00:00 2001
From: Teddy Koker <teddy.koker@gmail.com>
Date: Fri, 6 Mar 2026 09:21:27 -0500
Subject: [PATCH 2/2] version

---
 pyproject.toml | 2 +-
 uv.lock        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 75f634a..3449db4 100644
--- a/pyproject.toml
+++ b/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 = [
diff --git a/uv.lock b/uv.lock
index 7c18a6b..e2dab6b 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2233,7 +2233,7 @@ wheels = [
 
 [[package]]
 name = "nequix"
-version = "0.4.2"
+version = "0.4.3"
 source = { editable = "." }
 dependencies = [
     { name = "ase" },