fix: Fix keep_dims functionality for aten::max

[peri044]

Signed-off-by: dperi dperi@nvidia.com

Description

Fix keep_dims functionality for aten::max

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes

[github-actions]

There are some changes that do not conform to C++ style guidelines:

diff --git a/workspace/core/conversion/converters/impl/max.cpp b/tmp/changes.txt
index 693cd04..1f03b97 100644
--- a/workspace/core/conversion/converters/impl/max.cpp
+++ b/tmp/changes.txt
@@ -31,18 +31,18 @@ auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patter

       nvinfer1::ITensor* out0;
       nvinfer1::ITensor* out1;
-       if (!keep_dims){
+       if (!keep_dims) {
         if (topk_dims[dim] == 1) {
-             auto squeeze_layer = ctx->net->addShuffle(*topk_layer->getOutput(0));
-             squeeze_layer->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(0)->getDimensions(), dim));
-             TORCHTRT_CHECK(squeeze_layer, "Unable to create squeeze_layer layer from node: " << *n);
-             out0 = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer->getOutput(0));
-
-             auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
-             squeeze_layer_indices->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
-             TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
-             out1 = ctx->AssociateValueAndTensor(n->outputs()[1], squeeze_layer_indices->getOutput(0));
+           auto squeeze_layer = ctx->net->addShuffle(*topk_layer->getOutput(0));
+           squeeze_layer->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(0)->getDimensions(), dim));
+           TORCHTRT_CHECK(squeeze_layer, "Unable to create squeeze_layer layer from node: " << *n);
+           out0 = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer->getOutput(0));

+           auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
+           squeeze_layer_indices->setReshapeDimensions(
+               util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
+           TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
+           out1 = ctx->AssociateValueAndTensor(n->outputs()[1], squeeze_layer_indices->getOutput(0));
         }
       } else {
         out0 = ctx->AssociateValueAndTensor(n->outputs()[0], topk_layer->getOutput(0));
ERROR: Some files do not conform to style guidelines

[github-actions]

Code conforms to Python style guidelines

[narendasan]

Seems fine to me

[github-actions]

Code conforms to Python style guidelines

[github-actions]

Code conforms to C++ style guidelines