update comments and rebase

pytorch · Feb 9, 2024 · 8c41f14 · 8c41f14
1 parent 2bb9e2b
commit 8c41f14
Showing 1 changed file with 5 additions and 3 deletions.
diff --git 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
                 )