fix: fix the schema not found for node error

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

Author: bowang007

core/partitioning/partitioning.cpp

@@ -31,10 +31,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;
       }