diff --git a/tests/test_integration_sliding_window.py b/tests/test_integration_sliding_window.py
index c4d020276e..faec377586 100644
--- a/tests/test_integration_sliding_window.py
+++ b/tests/test_integration_sliding_window.py
@@ -84,7 +84,7 @@ def test_training(self):
             )
             output_image = nib.load(output_file).get_fdata()
             np.testing.assert_allclose(np.sum(output_image), 33621)
-            np.testing.assert_allclose(output_image.shape, (28, 25, 63, 1))
+            np.testing.assert_allclose(output_image.shape, (28, 25, 63))
 
 
 if __name__ == "__main__":