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

This commit is contained in:
Keyun Tong
2025-01-15 09:44:56 +00:00
committed by GitHub
parent cbe94391eb
commit ad388d25a8
2 changed files with 2 additions and 1 deletions
+1
View File
@@ -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):
+1 -1
View File
@@ -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