From 8c41f14fe5158559285c7d6001ba031bf3165ba1 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Thu, 8 Feb 2024 16:55:53 -0800 Subject: [PATCH] update comments and rebase --- py/torch_tensorrt/dynamo/conversion/impl/pool.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/pool.py b/py/torch_tensorrt/dynamo/conversion/impl/pool.py index e26e33ff64..5f5923c853 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/pool.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/pool.py @@ -136,7 +136,7 @@ def adaptive_avg_poolNd( axis = -axis positive_axis = get_positive_dim( axis, input_rank - ) # this is for calculating new shapes below + ) # convert to positive axis, which is for calculating new shapes below input_dim = input_shape[axis] output_dim = output_size[axis] diff = output_dim - input_dim @@ -156,9 +156,11 @@ def adaptive_avg_poolNd( remainder = 0 times += 1 - flags = [] + flags = [] # record the axis that needs to be repeated concat_list = [] - for j in range(input_dim): + for j in range( + input_dim + ): # iterate the input dim to see which dim needs to be repeated or not single_elem = impl.select.select( ctx, target, source_ir, f"{name}_select_{axis}_{j}", input, axis, j )