1 files changed, 2 insertions, 2 deletions
diff --git a/pseudocode/operators/ARGMAX.tosac b/pseudocode/operators/ARGMAX.tosac
index 982957a..73b74a0 100644
--- a/pseudocode/operators/ARGMAX.tosac
+++ b/pseudocode/operators/ARGMAX.tosac
@@ -24,14 +24,14 @@ for_each(left_index in left_shape) {
         in_t max_value = minimum_s<in_t>;
         out_t max_index = 0;
         for (i = 0; i < shape1[axis]; i++) {
-            dim_t index = flatten(left_index, [i], right_index);
+            shape_t index = flatten(left_index, [i], right_index);
             in_t value = tensor_read<in_t>(input, shape1, index);
             if (apply_max_s<in_t>(value, max_value) != max_value) {
                 max_value = value;
                 max_index = i;
             }
         }
-        dim_t index = flatten(left_index, right_index);
+        shape_t index = flatten(left_index, right_index);
         tensor_write<out_t>(output, shape, index, max_index);
     }
 }