facing issue while exporting smolvla_base model

[vijayrana-acl]

[vijayrana-acl]

[Bug] SmolVLA ONNX export fails with RuntimeError: expand: attempting to expand a dimension of length 0 -> 1024

Description
While attempting to convert the smolvla_base model into ONNX/TensorRT using reflex-vla on a Jetson Orin NX (16GB), the exporter crashes during the tensor broadcasting/shape expansion process.

The error originates deeply from PyTorch's fake_tensor tracer and ONNX diagnostics wrapper patches (patch_torch.py), throwing a RuntimeError due to an unexpected 0 dimension being expanded to 1024.

Environment Info
Hardware: NVIDIA Jetson Orin NX 16GB

VLA Model: smolvla_base

Stack Trace
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/torch/_refs/init.py", line 1129, in add
a, b = _maybe_broadcast(a, b)
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/onnx_diagnostic/torch_export_patches/patches/patch_torch.py", line 1029, in patched__maybe_broadcast
return tuple(_maybe_broadcast(x, common_shape) for x in args)
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/onnx_diagnostic/torch_export_patches/patches/patch_torch.py", line 1029, in
return tuple(_maybe_broadcast(x, common_shape) for x in args)
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/onnx_diagnostic/torch_export_patches/patches/patch_torch.py", line 1023, in __maybe_broadcast
return x.expand(common_shape)
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/torch/_refs/init.py", line 3042, in expand
torch._check(
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/torch/init.py", line 1690, in _check
_check_with(RuntimeError, cond, message)
File "/home/nvidia/Documents/vijay/groot_model_inf/n1_7/Isaac-GR00T/new_venv/lib/python3.10/site-packages/torch/init.py", line 1672, in _check_with
raise error_type(message_evaluated)
RuntimeError: expand: attempting to expand a dimension of length 0 -> 1024!
Export failed: RuntimeError: expand: attempting to expand a dimension of length 0 -> 1024!

To Reproduce / Steps to Trigger

1. Set up the virtual environment on Jetson Orin NX.

2. Run the export script/command for smolvla_base:
   (Add the specific CLI command or python snippet I used , reflex export --model smolvla_base ...)

3. The execution crashes at the torch.export patching phase.

[vijayrana-acl]

pip list of venv i used:

Package Version

---

absl-py 2.3.1
accelerate 1.13.0
action-msgs 1.2.2
action-tutorials-interfaces 0.20.9
action-tutorials-py 0.20.9
actionlib-msgs 4.9.1
aiofiles 25.1.0
aiohappyeyeballs 2.6.1
aiohttp 3.13.5
aiosignal 1.4.0
albucore 0.0.17
albumentations 1.4.18
ament-cmake-test 1.3.14
ament-copyright 0.12.15
ament-cppcheck 0.12.15
ament-cpplint 0.12.15
ament-flake8 0.12.15
ament-index-python 1.4.1
ament-lint 0.12.15
ament-lint-cmake 0.12.15
ament-package 0.14.2
ament-pep257 0.12.15
ament-uncrustify 0.12.15
ament-xmllint 0.12.15
angles 1.15.0
annotated-doc 0.0.4
annotated-types 0.7.0
antlr4-python3-runtime 4.9.3
anyio 4.13.0
array_record 0.8.1
astunparse 1.6.3
async-timeout 5.0.1
attrs 26.1.0
av 15.1.0
bitsandbytes 0.50.0.dev0
black 26.5.1
builtin-interfaces 1.2.2
certifi 2026.4.22
cffi 2.0.0
cfgv 3.5.0
charset-normalizer 3.4.7
chex 0.1.90
click 8.4.0
cloudpickle 3.1.2
cmake 3.31.10
coloredlogs 15.0.1
composition-interfaces 1.2.2
compressed-tensors 0.15.0.1
contourpy 1.3.2
control-msgs 4.9.0
controller-manager 2.54.0
controller-manager-msgs 2.54.0
cramjam 2.11.0
cryptography 42.0.8
cv-bridge 3.2.1
cycler 0.12.1
Cython 3.2.5
datasets 4.4.2
decorator 5.3.1
deepdiff 8.6.2
deepspeed 0.17.6
demo-nodes-py 0.20.9
diagnostic-msgs 4.9.1
diffusers 0.35.1
dill 0.4.0
distlib 0.4.0
distrax 0.1.5
dlimp 0.0.1
dm-tree 0.1.8
DockerHub-API 0.5
docopt 0.6.2
docstring_parser 0.18.0
domain-coordinator 0.10.0
draccus 0.10.0
einops 0.8.2
etils 1.13.0
eval_type_backport 0.3.1
evdev 1.9.3
example-interfaces 0.9.3
examples-rclpy-executors 0.15.5
examples-rclpy-minimal-action-client 0.15.5
examples-rclpy-minimal-action-server 0.15.5
examples-rclpy-minimal-client 0.15.5
examples-rclpy-minimal-publisher 0.15.5
examples-rclpy-minimal-service 0.15.5
examples-rclpy-minimal-subscriber 0.15.5
exceptiongroup 1.3.1
Farama-Notifications 0.0.6
fastparquet 2024.11.0
filelock 3.29.0
flake8 7.3.0
flatbuffers 25.12.19
flax 0.10.5
fonttools 4.63.0
frozenlist 1.8.0
fsspec 2025.10.0
furl 2.1.4
future 0.18.2
gast 0.4.0
generate-parameter-library-py 0.7.1
geometry-msgs 4.9.1
gitdb 4.0.12
GitPython 3.1.50
google-auth 2.53.0
google-auth-oauthlib 1.0.0
google-pasta 0.2.0
gr00t 0.1.0
graph-msgs 0.2.0
grpcio 1.80.0
gymnasium 1.3.0
h11 0.16.0
h5py 3.16.0
hf_transfer 0.1.9
hf-xet 1.5.0
hjson 3.1.0
httpcore 1.0.9
httpx 0.28.1
huggingface_hub 1.16.1
humanfriendly 10.0
humanize 4.15.0
identify 2.6.19
idna 3.15
image-geometry 3.2.1
ImageIO 2.37.3
imageio-ffmpeg 0.6.0
immutabledict 4.3.1
importlib_metadata 9.0.0
importlib_resources 7.1.0
inquirerpy 0.3.4
interactive-markers 2.3.2
jax 0.5.3
jax-cuda12-pjrt 0.5.3
jax-cuda12-plugin 0.5.3
jaxlib 0.5.3
Jinja2 3.1.6
joblib 1.5.3
joint-state-publisher 2.4.0
joint-state-publisher-gui 2.4.0
jsonlines 4.0.0
keras 3.12.2
Keras-Applications 1.0.8
Keras-Preprocessing 1.1.2
kiwisolver 1.5.0
laser-geometry 2.4.1
launch 1.0.14
launch-param-builder 0.1.1
launch-ros 0.19.13
launch-testing 1.0.14
launch-testing-ros 0.19.13
launch-xml 1.0.14
launch-yaml 1.0.14
lazy-loader 0.5
lerobot 0.4.4
libclang 18.1.1
lifecycle-msgs 1.2.2
lit 18.1.8
lmdb 1.7.5
logging-demo 0.20.9
loguru 0.7.3
map-msgs 2.1.0
Markdown 3.10.2
markdown-it-py 4.2.0
MarkupSafe 3.0.3
matplotlib 3.10.0
mccabe 0.7.0
mdurl 0.1.2
mergedeep 1.3.4
message-filters 4.3.16
micro-ros-msgs 1.0.0
ml_dtypes 0.5.0
mocap4r2-control-msgs 0.0.7
mock 3.0.5
moveit 0.1.3
moveit-configs-utils 2.5.9
moveit-msgs 2.2.1
moveit-task-constructor-msgs 0.1.3
mpmath 1.3.0
msgpack 1.1.0
msgpack-numpy 0.4.8
multidict 6.7.1
multiprocess 0.70.18
mypy_extensions 1.1.0
namex 0.1.0
nav-msgs 4.9.1
nest-asyncio 1.6.0
networkx 3.4.2
ninja 1.13.0
nodeenv 1.10.0
num2words 0.5.14
numpy 1.26.4
nvidia-cublas-cu12 12.9.2.10
nvidia-cuda-cccl-cu12 12.9.27
nvidia-cuda-cupti-cu12 12.9.79
nvidia-cuda-nvcc-cu12 12.9.86
nvidia-cuda-nvrtc-cu12 12.9.86
nvidia-cuda-runtime-cu12 12.9.79
nvidia-cudnn-cu12 9.22.0.52
nvidia-cufft-cu12 11.4.1.4
nvidia-cusolver-cu12 11.7.5.82
nvidia-cusparse-cu12 12.5.10.65
nvidia-ml-py 13.595.45
nvidia-modelopt 0.42.0
nvidia-nccl-cu12 2.30.4
nvidia-nvjitlink-cu12 12.9.86
nvidia-nvshmem-cu12 3.6.5
oauthlib 3.3.1
object-recognition-msgs 2.0.0
octomap-msgs 2.0.1
omegaconf 2.3.0
onnx 1.21.0
onnx-diagnostic 0.9.3
onnx-ir 0.2.1
onnxruntime-gpu 1.23.0
onnxscript 0.7.0
openarm_dataset 0.1.0
opencv-python 4.11.0.86
opencv-python-headless 4.11.0.86
opt_einsum 3.4.0
optax 0.2.8
optree 0.19.1
orbax-checkpoint 0.11.12
orderedmultidict 1.0.2
orderly-set 5.5.0
osrf-pycommon 2.1.6
packaging 25.0
pandas 2.3.3
pathspec 1.1.1
pcl-msgs 1.0.0
peft 0.17.1
pendulum-msgs 0.20.9
pfzy 0.3.4
pillow 12.2.0
pip 26.1.1
pkgconfig 1.6.0
platformdirs 4.9.6
play-motion-builder-msgs 1.4.1
play-motion2 1.8.4
play-motion2-msgs 1.8.4
pre_commit 4.6.0
profilehooks 1.13.0
promise 2.3
prompt_toolkit 3.0.52
propcache 0.5.2
protobuf 5.29.6
psutil 7.2.2
PuLP 3.3.1
py-cpuinfo 9.0.0
pyarrow 24.0.0
pyasn1 0.6.3
pyasn1_modules 0.4.2
pybind11 3.0.4
pycodestyle 2.14.0
pycparser 3.0
pydantic 2.13.4
pydantic_core 2.46.4
pyflakes 3.4.0
Pygments 2.20.0
pynput 1.8.2
pyparsing 3.3.2
pyserial 3.5
python-dateutil 2.9.0.post0
python-discovery 1.3.1
python-qt-binding 1.1.3
python-xlib 0.33
pytokens 0.4.1
pytz 2026.2
PyYAML 6.0.3
pyyaml-include 1.4.1
pyzmq 27.0.1
qt-dotgraph 2.2.5
qt-gui 2.2.5
qt-gui-cpp 2.2.5
qt-gui-py-common 2.2.5
quality-of-service-demo-py 0.20.9
qwen-vl-utils 0.0.14
rcl-interfaces 1.2.2
rclpy 3.3.21
rcutils 5.1.8
reflex-vla 0.9.6
regex 2026.5.9
requests 2.34.2
requests-oauthlib 2.0.0
rerun-sdk 0.26.2
resource-retriever 3.1.3
rich 15.0.0
rmw-dds-common 1.6.0
ros2action 0.18.18
ros2bag 0.15.16
ros2cli 0.18.18
ros2component 0.18.18
ros2doctor 0.18.18
ros2interface 0.18.18
ros2launch 0.19.13
ros2lifecycle 0.18.18
ros2multicast 0.18.18
ros2node 0.18.18
ros2param 0.18.18
ros2pkg 0.18.18
ros2plugin 5.1.4
ros2run 0.18.18
ros2service 0.18.18
ros2topic 0.18.18
rosbag2-interfaces 0.15.16
rosbag2-py 0.15.16
rosgraph-msgs 1.2.2
rosidl-adapter 3.1.8
rosidl-cli 3.1.8
rosidl-cmake 3.1.8
rosidl-generator-c 3.1.8
rosidl-generator-cpp 3.1.8
rosidl-generator-py 0.14.6
rosidl-generator-rs 0.4.12
rosidl-parser 3.1.8
rosidl-runtime-py 0.9.3
rosidl-typesupport-c 2.0.2
rosidl-typesupport-cpp 2.0.2
rosidl-typesupport-fastrtps-c 2.2.4
rosidl-typesupport-fastrtps-cpp 2.2.4
rosidl-typesupport-introspection-c 3.1.8
rosidl-typesupport-introspection-cpp 3.1.8
rpyutils 0.2.2
rqt 1.1.9
rqt-action 2.0.1
rqt-bag 1.1.6
rqt-bag-plugins 1.1.6
rqt-console 2.0.3
rqt-controller-manager 2.54.0
rqt-dotgraph 0.0.5
rqt-gauges 0.0.3
rqt-graph 1.3.2
rqt-gui 1.1.9
rqt-gui-py 1.1.9
rqt-joint-trajectory-controller 2.53.1
rqt-moveit 1.0.1
rqt-msg 1.2.0
rqt-plot 1.1.5
rqt-publisher 1.5.0
rqt-py-common 1.1.9
rqt-py-console 1.0.2
rqt-reconfigure 1.1.4
rqt-robot-dashboard 0.5.8
rqt-robot-monitor 1.0.6
rqt-robot-steering 1.0.3
rqt-runtime-monitor 1.0.0
rqt-service-caller 1.0.5
rqt-shell 1.0.2
rqt-srv 1.0.3
rqt-tf-tree 1.0.6
rqt-topic 1.5.1
safetensors 0.7.0
scikit-image 0.25.2
scikit-learn 1.7.2
scipy 1.15.3
sensor-msgs 4.9.1
sensor-msgs-py 4.9.1
sentry-sdk 2.60.0
setuptools 80.10.2
shape-msgs 4.9.1
shellingham 1.5.4
shtab 1.8.0
simple-parsing 0.1.8
simplejson 4.1.1
six 1.17.0
smmap 5.0.3
srdfdom 2.0.7
sros2 0.10.9
statistics-msgs 1.2.2
std-msgs 4.9.1
std-srvs 4.9.1
stereo-msgs 4.9.1
sympy 1.14.0
tabulate 0.10.0
teleop-twist-keyboard 2.4.1
tensorboard 2.18.0
tensorboard-data-server 0.7.2
tensorflow 2.18.0rc0
tensorflow-cpu-aws 2.15.1
tensorflow-datasets 4.9.10
tensorflow-estimator 2.15.0
tensorflow-io-gcs-filesystem 0.37.1
tensorflow-metadata 1.17.3
tensorflow-probability 0.25.0
tensorflow-text 2.18.1
tensorstore 0.1.78
termcolor 3.3.0
testresources 2.1.2
tf_keras 2.18.0
tf2-geometry-msgs 0.25.20
tf2-kdl 0.25.20
tf2-msgs 0.25.20
tf2-py 0.25.20
tf2-ros-py 0.25.20
tf2-sensor-msgs 0.25.20
tf2-tools 0.25.20
threadpoolctl 3.6.0
tifffile 2025.5.10
tokenizers 0.22.2
toml 0.10.2
tomli 2.4.1
toolz 1.1.0
topic-monitor 0.20.9
torch 2.8.0
torchcodec 0.7.0
torchvision 0.23.0
tqdm 4.67.3
trajectory-msgs 4.9.1
transformers 5.3.0
treescope 0.1.10
triton 3.5.1
turtlesim 1.4.3
typeguard 4.5.2
typer 0.25.1
typing_extensions 4.15.0
typing-inspect 0.9.0
typing-inspection 0.4.2
tyro 0.9.17
tzdata 2026.2
unique-identifier-msgs 2.2.1
urdfdom-py 1.2.1
urllib3 2.7.0
uvloop 0.22.1
virtualenv 21.3.3
vision-msgs 4.1.1
visualization-msgs 4.9.1
wandb 0.24.2
wcwidth 0.7.0
Werkzeug 3.1.8
wget 3.2
wheel 0.47.0
wrapt 1.14.2
xacro 2.1.1
xxhash 3.7.0
yarl 1.24.2
zipp 4.1.0

[DsThakurRawat]

hey you can start working on this
@rylinjames LGTM