diff --git a/ferminet/configs/excited/atoms.py b/ferminet/configs/excited/atoms.py
index bd6dd8a..14506f4 100644
--- a/ferminet/configs/excited/atoms.py
+++ b/ferminet/configs/excited/atoms.py
@@ -15,8 +15,8 @@
 """Config to reproduce Fig. 1 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs import atom
-from ferminet.jax.configs.excited import presets
+from ferminet.configs import atom
+from ferminet.configs.excited import presets
 import ml_collections
 
 
diff --git a/ferminet/configs/excited/benzene.py b/ferminet/configs/excited/benzene.py
index df5ed08..7bad38c 100644
--- a/ferminet/configs/excited/benzene.py
+++ b/ferminet/configs/excited/benzene.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 6 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 
diff --git a/ferminet/configs/excited/carbon_dimer.py b/ferminet/configs/excited/carbon_dimer.py
index 3969bd9..cea4593 100644
--- a/ferminet/configs/excited/carbon_dimer.py
+++ b/ferminet/configs/excited/carbon_dimer.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 3 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 
diff --git a/ferminet/configs/excited/double_excitation.py b/ferminet/configs/excited/double_excitation.py
index 1db3c87..cb81fe2 100644
--- a/ferminet/configs/excited/double_excitation.py
+++ b/ferminet/configs/excited/double_excitation.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 5 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 
diff --git a/ferminet/configs/excited/oscillator.py b/ferminet/configs/excited/oscillator.py
index 0a46f28..3ea4ca8 100644
--- a/ferminet/configs/excited/oscillator.py
+++ b/ferminet/configs/excited/oscillator.py
@@ -16,7 +16,7 @@
 
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 import pyscf
diff --git a/ferminet/configs/excited/twisted_ethylene.py b/ferminet/configs/excited/twisted_ethylene.py
index 376500b..90e0d5b 100644
--- a/ferminet/configs/excited/twisted_ethylene.py
+++ b/ferminet/configs/excited/twisted_ethylene.py
@@ -15,7 +15,7 @@
 """Config to reproduce Fig. 4 from Pfau et al. (2024)."""
 
 from ferminet import base_config
-from ferminet.jax.configs.excited import presets
+from ferminet.configs.excited import presets
 from ferminet.utils import system
 import ml_collections
 import numpy as np
diff --git a/ferminet/tests/excited_test.py b/ferminet/tests/excited_test.py
index a98a040..33febfe 100644
--- a/ferminet/tests/excited_test.py
+++ b/ferminet/tests/excited_test.py
@@ -23,13 +23,13 @@
 import chex
 from ferminet import base_config
 from ferminet import train
-from ferminet.jax.configs.excited import atoms
-from ferminet.jax.configs.excited import benzene
-from ferminet.jax.configs.excited import carbon_dimer
-from ferminet.jax.configs.excited import double_excitation
-from ferminet.jax.configs.excited import oscillator
-from ferminet.jax.configs.excited import presets
-from ferminet.jax.configs.excited import twisted_ethylene
+from ferminet.configs.excited import atoms
+from ferminet.configs.excited import benzene
+from ferminet.configs.excited import carbon_dimer
+from ferminet.configs.excited import double_excitation
+from ferminet.configs.excited import oscillator
+from ferminet.configs.excited import presets
+from ferminet.configs.excited import twisted_ethylene
 import jax
 
 import pyscf