From 4fec5bfa28572f5a829a4fb459e6dcb7cfa7a4e0 Mon Sep 17 00:00:00 2001 From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Date: Thu, 20 Feb 2025 15:17:24 +0800 Subject: [PATCH] Fix ONNX RuntimeError with dynamic WorldModel (#19322) Signed-off-by: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Co-authored-by: Glenn Jocher --- ultralytics/nn/tasks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 4d845ae9..5a5052fd 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -637,8 +637,8 @@ class WorldModel(DetectionModel): (torch.Tensor): Model's output tensor. """ txt_feats = (self.txt_feats if txt_feats is None else txt_feats).to(device=x.device, dtype=x.dtype) - if len(txt_feats) != len(x): - txt_feats = txt_feats.repeat(len(x), 1, 1) + if len(txt_feats) != len(x) or self.model[-1].export: + txt_feats = txt_feats.expand(x.shape[0], -1, -1) ori_txt_feats = txt_feats.clone() y, dt, embeddings = [], [], [] # outputs for m in self.model: # except the head part