Unverified Commit ad388d25 authored by Keyun Tong's avatar Keyun Tong Committed by GitHub
Browse files

Type-fix: make execute_model output type optional (#12020)

parent cbe94391
......@@ -70,6 +70,7 @@ class UniprocExecutor(Executor):
scheduler_output,
) -> ModelRunnerOutput:
output = self.worker.execute_model(scheduler_output)
assert output is not None
return output
def profile(self, is_start: bool = True):
......
......@@ -200,7 +200,7 @@ class Worker:
def execute_model(
self,
scheduler_output: "SchedulerOutput",
) -> ModelRunnerOutput:
) -> Optional[ModelRunnerOutput]:
output = self.model_runner.execute_model(scheduler_output)
return output if self.rank == 0 else None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment