diff --git a/.bumpversion.toml b/.bumpversion.toml
index 16a4a4d..9942a14 100644
--- a/.bumpversion.toml
+++ b/.bumpversion.toml
@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v1.1.0"
+current_version = "v1.1.1"
 commit = true
 commit_args = "--no-verify"
 tag = true
diff --git a/README.md b/README.md
index 2208e2c..cd49f25 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # invrs-gym
-`v1.1.0`
+`v1.1.1`
 
 ## Overview
 The `invrs_gym` package is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others.
diff --git a/pyproject.toml b/pyproject.toml
index 4c4c70f..466ee69 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_gym"
-version = "v1.1.0"
+version = "v1.1.1"
 description = "A collection of inverse design challenges"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"
diff --git a/src/invrs_gym/__init__.py b/src/invrs_gym/__init__.py
index 88e0f0d..e871c7e 100644
--- a/src/invrs_gym/__init__.py
+++ b/src/invrs_gym/__init__.py
@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v1.1.0"
+__version__ = "v1.1.1"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_gym import challenges as challenges
diff --git a/src/invrs_gym/challenges/diffract/splitter_challenge.py b/src/invrs_gym/challenges/diffract/splitter_challenge.py
index f913519..63bbecd 100644
--- a/src/invrs_gym/challenges/diffract/splitter_challenge.py
+++ b/src/invrs_gym/challenges/diffract/splitter_challenge.py
@@ -134,16 +134,16 @@ def eval_metric(self, response: common.GratingResponse) -> jnp.ndarray:
 
         # Total efficiency, i.e. sum of power into all output orders. A perfect
         # solution has a value of `1`, and the lowest possible value is `0`.
-        total_efficiency = jnp.sum(transmission, axis=(-2, -1))
+        total_efficiency = jnp.sum(transmission, axis=(-3, -2, -1))
 
         # Uniformity error. The perfect solution has a uniformity error of `0`,
         # and the highest possible value is `1`.
         uniformity_error = (
-            jnp.amax(transmission, axis=(-2, -1))
-            - jnp.amin(transmission, axis=(-2, -1))
+            jnp.amax(transmission, axis=(-3, -2, -1))
+            - jnp.amin(transmission, axis=(-3, -2, -1))
         ) / (
-            jnp.amax(transmission, axis=(-2, -1))
-            + jnp.amin(transmission, axis=(-2, -1))
+            jnp.amax(transmission, axis=(-3, -2, -1))
+            + jnp.amin(transmission, axis=(-3, -2, -1))
         )
 
         return jnp.amin(total_efficiency * (1 - uniformity_error))