From f628acadc5d2d0a9b093b8d9827e61b0fcc7315f Mon Sep 17 00:00:00 2001 From: Yiran Liu Date: Sat, 19 Mar 2022 12:26:13 +0800 Subject: [PATCH] feat: Add converter files for torch::max Signed-off-by: hongwei03 --- core/conversion/converters/impl/max.cpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp index 83f5aefd80..f9cd9f1954 100644 --- a/core/conversion/converters/impl/max.cpp +++ b/core/conversion/converters/impl/max.cpp @@ -17,17 +17,13 @@ auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patter {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { auto self = args[0].ITensorOrFreeze(ctx); - auto k = 1; auto dim = args[1].unwrapToInt(); - auto largest = true; auto selfDim = util::toVec(self->getDimensions()); if (dim < 0) { dim = selfDim.size() + dim; } uint32_t shiftDim = 1 << dim; - - auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN); - + auto TopKOperation = nvinfer1::TopKOperation::kMAX; auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim); TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n);