diff --git a/pygpe/shared/grid.py b/pygpe/shared/grid.py
index c4c9901..72b1ee7 100644
--- a/pygpe/shared/grid.py
+++ b/pygpe/shared/grid.py
@@ -127,7 +127,7 @@ def _generate_2d_grids(
             cp.arange(-self.num_points_y // 2, self.num_points_y // 2)
             * self.grid_spacing_y
         )
-        self.x_mesh, self.y_mesh = cp.meshgrid(x, y)
+        self.x_mesh, self.y_mesh = cp.meshgrid(x, y, indexing="ij")
 
         # Generate Fourier space variables
         self.fourier_spacing_x = cp.pi / (self.num_points_x // 2 * self.grid_spacing_x)
@@ -142,7 +142,9 @@ def _generate_2d_grids(
             * self.fourier_spacing_y
         )
 
-        self.fourier_x_mesh, self.fourier_y_mesh = cp.meshgrid(fourier_x, fourier_y)
+        self.fourier_x_mesh, self.fourier_y_mesh = cp.meshgrid(
+            fourier_x, fourier_y, indexing="ij"
+        )
         self.fourier_x_mesh = cp.fft.fftshift(self.fourier_x_mesh)
         self.fourier_y_mesh = cp.fft.fftshift(self.fourier_y_mesh)
 
@@ -181,7 +183,7 @@ def _generate_3d_grids(
             cp.arange(-self.num_points_z // 2, self.num_points_z // 2)
             * self.grid_spacing_z
         )
-        self.x_mesh, self.y_mesh, self.z_mesh = cp.meshgrid(x, y, z)
+        self.x_mesh, self.y_mesh, self.z_mesh = cp.meshgrid(x, y, z, indexing="ij")
 
         # Generate Fourier space variables
         self.fourier_spacing_x = cp.pi / (self.num_points_x // 2 * self.grid_spacing_x)
@@ -205,7 +207,7 @@ def _generate_3d_grids(
             self.fourier_x_mesh,
             self.fourier_y_mesh,
             self.fourier_z_mesh,
-        ) = cp.meshgrid(fourier_x, fourier_y, fourier_z)
+        ) = cp.meshgrid(fourier_x, fourier_y, fourier_z, indexing="ij")
         self.fourier_x_mesh = cp.fft.fftshift(self.fourier_x_mesh)
         self.fourier_y_mesh = cp.fft.fftshift(self.fourier_y_mesh)
         self.fourier_z_mesh = cp.fft.fftshift(self.fourier_z_mesh)
diff --git a/pyproject.toml b/pyproject.toml
index ab3aa04..ab6e3bb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pygpe"
-version = "2.0.0"
+version = "2.0.1"
 description = "A fast Gross-Pitaevskii equation solver for scalar, spin-1 and spin-2 BEC systems."
 authors = ["Matt Wheeler <wheelermt@proton.me>"]
 license = "MIT"
@@ -20,6 +20,7 @@ gpu = ["cupy"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.4"
+black = "^24.3.0"
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/tests/shared/test_grids.py b/tests/shared/test_grids.py
index fb89bf8..3e7b9af 100644
--- a/tests/shared/test_grids.py
+++ b/tests/shared/test_grids.py
@@ -59,9 +59,9 @@ def test_correct_fourier_shift_2d():
 
     grid2d = grid.Grid((64, 64), (0.5, 0.5))
 
-    for element in grid2d.fourier_x_mesh[:, 0]:
+    for element in grid2d.fourier_x_mesh[0, :]:
         assert element == 0
-    for element in grid2d.fourier_y_mesh[0, :]:
+    for element in grid2d.fourier_y_mesh[:, 0]:
         assert element == 0