In case someone stumble on this, the issue was actually with flexible input dimension, same model with fixed input shapes is behaving as expected

Have you filed a bug for this? no Have you tested on macOS Sequoia? the behaviour is reproducible on macOS Sequoia