Correcting the padding of mask tensor when the dimension is lesser th…

…an self tensor
pytorch · May 19, 2023 · 3b10531 · 3b10531
1 parent 98c4011
commit 3b10531
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp
@@ -708,12 +708,12 @@ auto select_registrations TORCHTRT_UNUSED =
                  int diff = mask->getDimensions().nbDims - self->getDimensions().nbDims;
                  for (int i = 0; i < diff; i++) {
                    if (mask->getDimensions().d[i] != 1) {
-                     LOG_DEBUG("The mask dimension cannot be unpadded to the self dimension");
+                     LOG_ERROR("The mask dimension cannot be unpadded to the self dimension");
                    }
                  }
                  mask = addUnpadding(ctx, n, mask, self->getDimensions().nbDims, false, true);
                } else {
-                 self = addPadding(ctx, n, self, mask->getDimensions().nbDims, false, true);
+                 mask = addPadding(ctx, n, mask, self->getDimensions().nbDims, false, true);
                }
                // mask = addPadding(ctx, n, mask, self->getDimensions().nbDims, false, true);
                auto val = args[2].unwrapToScalar();