resolve reviews

pytorch · Sep 22, 2023 · 9e62066 · 9e62066
1 parent 0047b3d
commit 9e62066
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 35 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1395,5 +1395,5 @@ def aten_ops_argmax(
         name,
         input=args[0],
         dim=args_bounds_check(args, 1),
-        keep_dim=args_bounds_check(args, 2),
+        keep_dim=args_bounds_check(args, 2, False),
     )
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/argmax.py b/py/torch_tensorrt/dynamo/conversion/impl/argmax.py
@@ -3,8 +3,14 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor
-from torch_tensorrt.fx.converters.converter_utils import set_layer_name
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    cast_trt_tensor,
+    get_axes_for_reduce_op,
+)
+from torch_tensorrt.fx.converters.converter_utils import (
+    get_positive_dim,
+    set_layer_name,
+)
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
 from . import squeeze
@@ -27,7 +33,7 @@ def argmax(
         input = cast_trt_tensor(network, input, trt.float32, name)
     if dim < 0:
         dim = len(tuple(input.shape)) + dim
-    reduce_mask = 1 << dim
+    reduce_mask = get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape)))
     topk_layer = network.add_topk(input, trt.TopKOperation.MAX, 1, reduce_mask)
     set_layer_name(topk_layer, target, name)
 

diff --git a/tests/py/dynamo/converters/test_argmax.py b/tests/py/dynamo/converters/test_argmax.py