[BugFix] Fix assuming all stage model have talker (#730)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
汪志鹏 authored and GitHub committed 2026-01-11 01:36:37 +13:00
1 parent f3741ac2b9
commit 0ec2da08de
1 file changed
+5 -1
+5 -1
View File
@@ -563,7 +563,11 @@ class OmniGPUModelRunner(GPUModelRunner):
ubatch_slices=ubatch_slices,
),
):
if self.model.talker is not None and hasattr(self.model, "talker_mtp") and num_tokens_padded == 1:
if (
getattr(self.model, "talker", None) is not None
and hasattr(self.model, "talker_mtp")
and num_tokens_padded == 1
):
outputs = self.talker_mtp(
self.talker_mtp_input_ids.gpu[:num_tokens_padded],
self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded],