diff --git a/src/sagemaker/image_uri_config/neo-mxnet.json b/src/sagemaker/image_uri_config/neo-mxnet.json
index d353c5830a..55cc9a022b 100644
--- a/src/sagemaker/image_uri_config/neo-mxnet.json
+++ b/src/sagemaker/image_uri_config/neo-mxnet.json
@@ -2,20 +2,21 @@
     "processors": ["cpu", "gpu"],
     "scope": ["inference"],
     "version_aliases": {
-        "0.12.1": "1.7",
-        "1.0.0": "1.7",
-        "1.1.0": "1.7",
-        "1.2": "1.7",
-        "1.2.0": "1.7",
-        "1.2.1": "1.7",
-        "1.3": "1.7",
-        "1.3.0": "1.7",
-        "1.4": "1.7",
-        "1.4.0": "1.7",
-        "1.4.1": "1.7"
+        "0.12.1": "1.8",
+        "1.0.0": "1.8",
+        "1.1.0": "1.8",
+        "1.2": "1.8",
+        "1.2.0": "1.8",
+        "1.2.1": "1.8",
+        "1.3": "1.8",
+        "1.3.0": "1.8",
+        "1.4": "1.8",
+        "1.4.0": "1.8",
+        "1.4.1": "1.8",
+        "1.7": "1.8"
     },
     "versions": {
-        "1.7": {
+        "1.8": {
             "py_versions": ["py3"],
             "registries": {
                 "af-south-1": "774647643957",
diff --git a/tests/data/mxnet_mnist/mnist_neo.py b/tests/data/mxnet_mnist/mnist_neo.py
index 1698a49154..ede39af4c9 100644
--- a/tests/data/mxnet_mnist/mnist_neo.py
+++ b/tests/data/mxnet_mnist/mnist_neo.py
@@ -105,7 +105,7 @@ def train(
 
 
 def model_fn(path_to_model_files):
-    import neomxnet  # noqa: F401
+    import neomx  # noqa: F401
 
     ctx = mx.cpu()
     sym, arg_params, aux_params = mx.model.load_checkpoint(
@@ -120,7 +120,7 @@ def model_fn(path_to_model_files):
 
 
 def transform_fn(mod, payload, input_content_type, requested_output_content_type):
-    import neomxnet  # noqa: F401
+    import neomx  # noqa: F401
 
     if input_content_type != "application/vnd+python.numpy+binary":
         raise RuntimeError("Input content type must be application/vnd+python.numpy+binary")
diff --git a/tests/integ/test_neo_mxnet.py b/tests/integ/test_neo_mxnet.py
index 3237b36e3e..61fc59b0e2 100644
--- a/tests/integ/test_neo_mxnet.py
+++ b/tests/integ/test_neo_mxnet.py
@@ -71,7 +71,7 @@ def test_attach_deploy(
 
         estimator.compile_model(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             output_path=estimator.output_path,
         )
 
@@ -121,7 +121,7 @@ def test_deploy_model(
 
         model.compile(
             target_instance_family=cpu_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),
@@ -163,7 +163,7 @@ def test_inferentia_deploy_model(
 
         model.compile(
             target_instance_family=inf_instance_family,
-            input_shape={"data": [1, 1, 28, 28], "softmax_label": [1]},
+            input_shape={"data": [1, 1, 28, 28]},
             role=role,
             job_name=unique_name_from_base("test-deploy-model-compilation-job"),
             output_path="/".join(model_data.split("/")[:-1]),