From 208d2e15a9b824f56583f202b4ee206138d94082 Mon Sep 17 00:00:00 2001 From: ice-tong Date: Thu, 22 Aug 2024 18:54:32 +0800 Subject: [PATCH] Use torch.amax in ReduceMax forward to support multiplie dim --- ppq/executor/op/torch/default.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ppq/executor/op/torch/default.py b/ppq/executor/op/torch/default.py index db93e773..a671beac 100644 --- a/ppq/executor/op/torch/default.py +++ b/ppq/executor/op/torch/default.py @@ -1619,7 +1619,7 @@ def ReduceMax_forward(op: Operation, values: List[torch.Tensor], ctx: TorchBacke if keepdim: output = output.reshape([1] * input_value.dim()) else: - output, _ = torch.max(input_value, dim=dim[0], keepdim=keepdim) + output = torch.amax(input_value, dim=dim, keepdim=keepdim) return output