feat(aten::add): adding string concat evaluator

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · May 8, 2022 · 65dbf90 · 65dbf90
1 parent 828d120
commit 65dbf90
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp
@@ -342,15 +342,22 @@ auto aten_registrations TORCHTRT_UNUSED =
                         auto a = args.at(n->input(0)).unwrapToDouble();
                         auto b = args.at(n->input(1)).unwrapToDouble();
                         return a + b;
+                      } else if (args.at(n->input(0)).IValue()->isString()) {
+                        auto a = args.at(n->input(0)).unwrapToString();
+                        auto b = args.at(n->input(1)).unwrapToString();
+                        return a + b;
                       } else {
                         TORCHTRT_THROW_ERROR(
                             "Unimplemented data type for aten::add evaluator: "
                             << args.at(n->input(0)).IValue()->type()->str());
                         return {};
                       }
                     },
-                    EvalOptions().validSchemas(
-                        {"aten::add.int(int a, int b) -> (int)", "aten::add.float(float a, float b) -> (float)"})})
+                    EvalOptions().validSchemas({
+                      "aten::add.int(int a, int b) -> (int)",
+                      "aten::add.float(float a, float b) -> (float)",
+                      "aten::add.str(str a, str b) -> (str)"
+                    })})
         .evaluator({c10::Symbol::fromQualString("aten::add_"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       if (args.at(n->input(0)).IValue()->isList()) {