Fix: Shape mismatch in extend mode causing AssertionError#373

ChuxiJ merged 1 commit intoace-step:mainace-step/ACE-Step:mainfrom

iackov:fix/extend-mode-shape-mismatchiackov/ACE-Step:fix/extend-mode-shape-mismatchCopy head branch name to clipboard