diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py
index b6f09ff..efc2752 100644
--- a/src/loch/_kernels.py
+++ b/src/loch/_kernels.py
@@ -498,7 +498,7 @@
                 water[2] = 0.0f;
 
                 // Shift the other atoms by the appropriate amount.
-                for (int i = 0; i < 3; i++)
+                for (int i = 0; i < num_points; i++)
                 {
                     water[i*3 + 0] = water_template[i*3 + 0] - water_template[0];
                     water[i*3 + 1] = water_template[i*3 + 1] - water_template[1];
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index 4f8574b..b8c9a95 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -2025,7 +2025,7 @@ def _accept_insertion(
 
         # Get the new water positions.
         water_positions = self._water_positions.get().reshape(
-            (self._batch_size, 3, self._num_points)
+            (self._batch_size, self._num_points, 3)
         )[idx]
 
         # Update the water state.