feat: Add converter files for torch::max

Signed-off-by: hongwei03 <[email protected]>
pytorch · Mar 23, 2022 · dd7a44e · dd7a44e
1 parent 0b5673b
commit dd7a44e
Showing 1 changed file with 18 additions and 18 deletions.
diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp
@@ -14,27 +14,27 @@ namespace converters {
 namespace impl {
 namespace {
 auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
-        {"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 dim = args[1].unwrapToInt();
-             auto selfDim = util::toVec(self->getDimensions());
-             if (dim < 0) {
-                 dim = selfDim.size() + dim;
-             }
-             uint32_t shiftDim = 1 << dim;
-             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);
+    {"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 dim = args[1].unwrapToInt();
+       auto selfDim = util::toVec(self->getDimensions());
+       if (dim < 0) {
+         dim = selfDim.size() + dim;
+       }
+       uint32_t shiftDim = 1 << dim;
+       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);
 
-             auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
-             auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1));
+       auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+       auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1));
 
-             LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
-             LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
+       LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
+       LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
 
-             return true;
-         }});
+       return true;
+     }});
 } // namespace
 } // namespace impl
 } // namespace converters