Skip to content

Commit f628aca

Browse files
Yiran Liup517332051
authored andcommitted
feat: Add converter files for torch::max
Signed-off-by: hongwei03 <[email protected]>
1 parent 569bcde commit f628aca

File tree

1 file changed

+1
-5
lines changed
  • core/conversion/converters/impl

1 file changed

+1
-5
lines changed

core/conversion/converters/impl/max.cpp

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,13 @@ auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patter
1717
{"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)",
1818
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
1919
auto self = args[0].ITensorOrFreeze(ctx);
20-
auto k = 1;
2120
auto dim = args[1].unwrapToInt();
22-
auto largest = true;
2321
auto selfDim = util::toVec(self->getDimensions());
2422
if (dim < 0) {
2523
dim = selfDim.size() + dim;
2624
}
2725
uint32_t shiftDim = 1 << dim;
28-
29-
auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN);
30-
26+
auto TopKOperation = nvinfer1::TopKOperation::kMAX;
3127
auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim);
3228
TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n);
3329

0 commit comments

Comments
 (0)