diff --git a/src/main/resources/templates/gluon/CNNSupervisedTrainer.ftl b/src/main/resources/templates/gluon/CNNSupervisedTrainer.ftl
index e47a002301975b7b79d999eb12012c3d0c3dd62a..34951554ad820535b13edb214e1b4b98afb8eee9 100644
--- a/src/main/resources/templates/gluon/CNNSupervisedTrainer.ftl
+++ b/src/main/resources/templates/gluon/CNNSupervisedTrainer.ftl
@@ -393,11 +393,7 @@ class ${tc.fileNameWithoutEnding}:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_Alexnet.py b/src/test/resources/target_code/CNNSupervisedTrainer_Alexnet.py
index 2f17688bd4e481792071b45dd78a0d88d57b15c4..8ad298bbb8006a370248bdfca9492d7fb7ca8418 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_Alexnet.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_Alexnet.py
@@ -512,11 +512,7 @@ class CNNSupervisedTrainer_Alexnet:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_CifarClassifierNetwork.py b/src/test/resources/target_code/CNNSupervisedTrainer_CifarClassifierNetwork.py
index 54f4c69af69f6fa5f68bf9ab0cc81e98307e9148..a53642cd4eb9cba837c557ecc8f0bcbc9d8df8c2 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_CifarClassifierNetwork.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_CifarClassifierNetwork.py
@@ -512,11 +512,7 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_Invariant.py b/src/test/resources/target_code/CNNSupervisedTrainer_Invariant.py
index cc63d68919f623789628b31ff137922f303bfe93..1eaeb70a80e79836d3ed4fbb23839d49e6f25cd9 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_Invariant.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_Invariant.py
@@ -497,11 +497,7 @@ class CNNSupervisedTrainer_Invariant:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_MultipleStreams.py b/src/test/resources/target_code/CNNSupervisedTrainer_MultipleStreams.py
index 0903c286d7d289c8bbb24bb291213a7764c4aaec..769883a4b928132f2d8604e12245d0b7aad1e8cc 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_MultipleStreams.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_MultipleStreams.py
@@ -485,11 +485,7 @@ class CNNSupervisedTrainer_MultipleStreams:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_RNNencdec.py b/src/test/resources/target_code/CNNSupervisedTrainer_RNNencdec.py
index 53765c13399c2bfaa684cf9a82d0452afb3a076d..1be1aeceff712f04ef18176c32b46d0e85faa3ad 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_RNNencdec.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_RNNencdec.py
@@ -599,11 +599,7 @@ class CNNSupervisedTrainer_RNNencdec:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_RNNsearch.py b/src/test/resources/target_code/CNNSupervisedTrainer_RNNsearch.py
index 068ca8c466d1934a8ed05b950fed6917481e5cae..cd1839bb2b93b9fd4185977120bc102741a5f7e0 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_RNNsearch.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_RNNsearch.py
@@ -596,11 +596,7 @@ class CNNSupervisedTrainer_RNNsearch:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_RNNtest.py b/src/test/resources/target_code/CNNSupervisedTrainer_RNNtest.py
index 052800cbb9bc67eafc94abc0d9503bde74b8055a..576bb0718d9ece21ea84aaf9b25ab6b6e39ae485 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_RNNtest.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_RNNtest.py
@@ -565,11 +565,7 @@ class CNNSupervisedTrainer_RNNtest:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_ResNeXt50.py b/src/test/resources/target_code/CNNSupervisedTrainer_ResNeXt50.py
index 957ae83997a7177262691a59318cdb977ff76106..7c50965d9c5b58622447b7aa3974d76b2e2d6cee 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_ResNeXt50.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_ResNeXt50.py
@@ -470,11 +470,7 @@ class CNNSupervisedTrainer_ResNeXt50:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_Show_attend_tell.py b/src/test/resources/target_code/CNNSupervisedTrainer_Show_attend_tell.py
index a223bc409a85a0d95c66b6244ec272124df2c28e..78f7d2002666c8c8a8e74ad3309e5fb28dad4a9e 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_Show_attend_tell.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_Show_attend_tell.py
@@ -587,11 +587,7 @@ class CNNSupervisedTrainer_Show_attend_tell:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_ThreeInputCNN_M14.py b/src/test/resources/target_code/CNNSupervisedTrainer_ThreeInputCNN_M14.py
index 7ed3647dc11e8e822364e8b94aae163421438b67..49fd0cd5e27d7ec3e10a444104328ecffb57f94c 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_ThreeInputCNN_M14.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_ThreeInputCNN_M14.py
@@ -476,11 +476,7 @@ class CNNSupervisedTrainer_ThreeInputCNN_M14:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]
diff --git a/src/test/resources/target_code/CNNSupervisedTrainer_VGG16.py b/src/test/resources/target_code/CNNSupervisedTrainer_VGG16.py
index 7f4d8a87e3e11d9950346eed91790d292c4f6a28..bcf5d85d3292f3a3608eef1ef343aadbd9c686b2 100644
--- a/src/test/resources/target_code/CNNSupervisedTrainer_VGG16.py
+++ b/src/test/resources/target_code/CNNSupervisedTrainer_VGG16.py
@@ -512,11 +512,7 @@ class CNNSupervisedTrainer_VGG16:
 
                 predictions = []
                 for output_name in outputs:
-                    if mx.nd.shape_array(mx.nd.squeeze(output_name)).size > 1:
-                        predictions.append(mx.nd.argmax(output_name, axis=1))
-                    #ArgMax already applied
-                    else:
-                        predictions.append(output_name)
+                    predictions.append(output_name)
 
                 metric.update(preds=predictions, labels=labels)
             test_metric_score = metric.get()[1]