chore: play around with aten::size

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

Author: peri044

core/conversion/converters/impl/shuffle.cpp

@@ -73,8 +73,9 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
                std::cout << "====2====" << std::endl;
                std::vector<int64_t> new_shape;
                if (ctx->input_is_dynamic) {
-                 std::cout << "====3====" << std::endl;
-                 new_shape = util::toVec(args[1].unwrapToIntList().vec());
+                 std::cout << "====3====: " << args[1].size() << std::endl;
+                 // new_shape = util::toVec(args[1].unwrapToIntList().vec());
+                 new_shape = util::toVec(args[1].unwrapToITensorList());
                  std::cout << "====4====" << std::endl;
                  int nbDynamicDims = 0;
                  for (size_t i = 0; i < new_shape.size(); i++) {

core/conversion/var/Var.cpp

@@ -146,6 +146,24 @@ bool Var::isITensor() const {
   }
 }
 
+bool Var::isITensorList() const {
+  LOG_DEBUG("===== TYPE NAME: " << type_name());
+  if (type_ == Type::kITensor) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+bool Var::unwrapToITensorList() {
+  TORCHTRT_CHECK(
+      isIValue(), "Requested unwrapping of arg assuming it was an IValue, however arg type is " << type_name());
+  LOG_DEBUG("===== TYPE NAME: " << type_name());
+  auto ivalue = ptr_.ivalue;
+  return false;
+  // return ptr_.ivalue->to<nvinfer1::ITensor*>();
+}
+
 bool Var::isIValue() const {
   if (type_ == Type::kIValue) {
     return true;

core/conversion/var/Var.h

@@ -43,6 +43,7 @@ class Var : torch::CustomClassHolder {
   c10::Scalar unwrapToScalar();
   c10::List<int64_t> unwrapToIntList(c10::List<int64_t> default_val);
   c10::List<int64_t> unwrapToIntList();
+  c10::List<nvinfer1::ITensor*> unwrapToITensorList();
   c10::List<double> unwrapToDoubleList(c10::List<double> default_val);
   c10::List<double> unwrapToDoubleList();
   c10::List<bool> unwrapToBoolList(c10::List<bool> default_val);
@@ -58,6 +59,7 @@ class Var : torch::CustomClassHolder {
 
   bool isIValue() const;
   bool isITensor() const;
+  bool isITensorList() const;
   bool isNone() const;
   Var::Type type() const;
   std::string type_name() const;