Merge pull request #2 from OpenBioSim/fix_1

Fix issue #1

Author: lohedges

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];

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.