diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 6be9dc2..6771710 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -134,7 +134,7 @@
     Add = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
     AddN = OperatorInfo()
     Any = OperatorInfo()
-    ArgMax = OperatorInfo()
+    ArgMax = OperatorInfo(indices=NNG_IFM_INDICES)
     ArgMin = OperatorInfo()
     AvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
     Atan2 = OperatorInfo(indices=NNG_IFM_IFM2_INDICES)