fix python segfault 2

flexflow · Jul 18, 2023 · 1367f8a · 1367f8a
1 parent 716a5b1
commit 1367f8a
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/runtime/model.cc b/src/runtime/model.cc
@@ -6002,8 +6002,15 @@ void register_flexflow_internal_tasks(Runtime *runtime,
                                    "AllReduce Inference");
     registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC));
     registrar.set_leaf();
-    Runtime::preregister_task_variant<AllReduce::inference_task>(
-        registrar, "AllReduce Inference Task");
+    if (pre_register) {
+      Runtime::preregister_task_variant<AllReduce::inference_task>(
+          registrar, "AllReduce Inference Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime->register_task_variant<AllReduce::inference_task>(registrar);
+    }
   }
   {
     TaskVariantRegistrar registrar(ALLREDUCE_FWD_TASK_ID, "AllReduce Forward");