fix: fix the schema not found for node error

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

core/partitioning/partitioning.cpp

@@ -40,10 +40,13 @@ bool containNonTensorOutputs(torch::jit::Node* n) {
 }
 
 bool isModifyingNodes(torch::jit::Node* node, torch::jit::Value* val) {
-  const auto& schema = node->schema();
+  const torch::jit::FunctionSchema* schema = node->maybeSchema();
+  if (!schema) {
+    return false;
+  }
   for (size_t i = 0; i < node->inputs().size(); ++i) {
     if (node->inputs()[i] == val) {
-      const at::AliasInfo* formal = schema.arguments()[i].alias_info();
+      const at::AliasInfo* formal = schema->arguments()[i].alias_info();
       if (formal && formal->isWrite()) {
         return true;
       }
@@ -124,7 +127,8 @@ void find_all_fallback_nodes(
       if (!isTensor(output)) {
         for (auto use : output->uses()) {
           auto node = use.user;
-          if (node->kind() != torch::jit::prim::Constant && global_fallback_nodes.insert({node, FallbackNodeType::kNON_TENSOR}).second) {
+          if (node->kind() != torch::jit::prim::Constant &&
+              global_fallback_nodes.insert({node, FallbackNodeType::kNON_TENSOR}).second) {
             q.push(node);
           }
         }