fix: Return static size for desired dimension if it's available

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/evaluators/aten.cpp

@@ -299,20 +299,22 @@ auto aten_registrations TORCHTRT_UNUSED =
                } else {
                  auto dim = args.at(n->input(1)).unwrapToInt();
                  if (tensor_var.isITensor()) {
-                   if (ctx->input_is_dynamic) {
+                   auto tensor = tensor_var.ITensor();
+                   auto dims = util::toVec(tensor->getDimensions());
+                   auto nbDims = tensor->getDimensions().nbDims;
+                   if (dim < 0) {
+                     dim += nbDims;
+                   }
+                   // Check if selected dimension size is -1 else return static size
+                   if (ctx->input_is_dynamic && dims[dim] == -1) {
                      if (ctx->settings.allow_shape_tensors) {
                        return dynamic_size_layer(ctx, n, args);
                      } else {
                        LOG_WARNING(
                            "There may be undefined behavior using dynamic shape and aten::size without setting allow_shape_tensors");
                      }
                    }
-                   auto tensor = tensor_var.ITensor();
-                   auto dims = util::toVec(tensor->getDimensions());
-                   auto nbDims = tensor->getDimensions().nbDims;
-                   if (dim < 0) {
-                     dim += nbDims;
-                   }
+
                    return dims[dim];
                  } else if (tensor_var.IValue()->isTensor()) {
                    auto tensor = tensor_var.unwrapToTensor();

core/conversion/evaluators/eval_util.cpp

@@ -45,6 +45,10 @@ c10::IValue dynamic_size_layer(ConversionCtx* ctx, const torch::jit::Node* n, kw
     // Handle negative axis by refering to nbDims of input Tensor
     dim = dim < 0 ? dim + maxDim : dim;
     LOG_DEBUG("Dimension to select: " << dim);
+    // Check if selected dimension size is -1 else return static size
+    if (input_dims.d[dim] != -1) {
+      return input_dims.d[dim];
+    }
     shape_1d_tensor = index_layer(ctx, n, shape_1d_tensor, dim);
     LOG_DEBUG("Output tensor shape: " << shape_1d_tensor->getDimensions());