VITRA VLA inference runs on AMD GPUs (ROCm) — verified on MI300X, zero source changes

Summary

VITRA's VLA inference path (build_vla + predict_action) runs on AMD Instinct GPUs via ROCm with no source changes beyond the dtype handling already documented in the README minimal example.

VITRA is pure PyTorch / transformers (no custom CUDA kernels), so it works on ROCm out of the box once the environment is set up correctly.

Tested on

• GPU: AMD Instinct MI300X (also verified on MI308X, gfx942)
• ROCm: 7.2.1
• PyTorch: 2.9.1 (ROCm build)
• Docker: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1
• transformers==4.47.1 (the repo's pin)
• Weights: VITRA-VLA/VITRA-VLA-3B + google/paligemma2-3b-mix-224

Result

• VITRA_Paligemma loads the checkpoint with 0 missing / 0 unexpected state-dict keys.
• predict_action returns a finite [1, 16, 192] action chunk (16-step × 192-dim unified action).
• Timings (warm cache): load ≈ 15.5s, infer ≈ 10.0s.

Setup notes for ROCm / AMD users

A few environment details that make it work — none require source changes:

1. Pin transformers==4.47.1 (the repo's pyproject pin). Newer transformers change
   get_image_features' return type → 'BaseModelOutputWithPooling' has no attribute 'numel',
   and the correct pin also makes the checkpoint load cleanly (0/0 keys).
2. Install without training-only CUDA deps for inference: a full pip install -e . stalls
   building flash-attn / deepspeed, which are only needed for training. Using
   pip install --no-deps -e . plus the explicit pure-Python deps keeps the Docker ROCm
   PyTorch intact (torch.version.hip unchanged) and inference runs fine without them.
   Would a pip install-friendly inference-only path (moving flash-attn/deepspeed to an
   optional extra) be welcome? Happy to help.
3. float32 inputs: state / fov built from NumPy / np.deg2rad default to float64
   (torch.Double) and raise mat1 and mat2 must have the same dtype. This is already handled
   by the README minimal example (which casts to float32), so no change needed — just noting
   it for other ROCm users following along.

Ask

Nothing blocking — this is mainly a compatibility report so AMD/ROCm users know VITRA works.
If useful, I'm happy to open a follow-up PR for the inference-only optional-extras packaging
in note (2).

Read the README minimal example and the repository's pyproject dependency pin, then compare them with the reported ROCm setup and dtype requirements. Done would mean a clearly scoped documentation update for AMD/ROCm inference; the optional inference-only packaging idea is presented as a possible follow-up rather than a defined change.