scripts/train_pytorch.py debug crashes: image layout conversion is gated on uint8, and the vision projection width is hardcoded

The command documented in the README under "PyTorch Training" fails on main:

uv run scripts/train_pytorch.py debug --exp_name pytorch_test

Two independent bugs cause this. Both reproduce on CPU, so no GPU or CUDA kernels are involved.

The JAX path is unaffected: uv run scripts/train.py debug --exp_name jax_debug_check completes 10 steps and checkpoints normally on the same commit and machine. Only the PyTorch path fails.

Environment

• OS: Ubuntu 24.04.4 LTS, x86_64
• Python 3.11.16
• openpi 215abfb (current main), unmodified
• transformers_replace copied into site-packages per the README
• CUDA_VISIBLE_DEVICES="" JAX_PLATFORMS=cpu

Bug 1: NHWC to NCHW conversion is gated on dtype == torch.uint8

Observation.from_dict performs the layout conversion only inside the uint8 branch:

https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/models/model.py#L116-L120

for key in data["image"]:
    if data["image"][key].dtype == np.uint8:
        data["image"][key] = data["image"][key].astype(np.float32) / 255.0 * 2.0 - 1.0
    elif hasattr(data["image"][key], "dtype") and data["image"][key].dtype == torch.uint8:
        data["image"][key] = data["image"][key].to(torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0

Normalization and layout conversion are orthogonal concerns sharing one condition. Real datasets deliver uint8, so both happen together and the path works. FakeDataset builds samples from Pi0Config.inputs_spec(), where images are float32 with shape [b, 224, 224, 3]:

https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/training/data_loader.py#L112-L113 https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/models/pi0_config.py#L65

Float32 images therefore skip the permute and reach SigLIP still in NHWC.

preprocess_observation_pytorch does not correct this. It sniffs the incoming layout and only restores channels-first if the input was already channels-first:

https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/models_pytorch/preprocessing_pytorch.py#L40-L45

# TODO: This is a hack to handle both [B, C, H, W] and [B, H, W, C] formats
is_channels_first = image.shape[1] == 3

For (2, 224, 224, 3), shape[1] is 224, so the tensor passes through unchanged.

Traceback

  File "scripts/train_pytorch.py", line 529, in train_loop
    losses = model(observation, actions)
  File "src/openpi/models_pytorch/pi0_pytorch.py", line 331, in forward
    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)
  File "src/openpi/models_pytorch/pi0_pytorch.py", line 203, in embed_prefix
    img_emb = self._apply_checkpoint(image_embed_func, img)
  ...
  File "transformers/models/siglip/modeling_siglip.py", line 274, in forward
    patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
  File "torch/nn/modules/conv.py", line 549, in _conv_forward
    return F.conv2d(
RuntimeError: Given groups=1, weight of size [1152, 3, 14, 14], expected input[2, 224, 224, 3] to have 3 channels, but got 224 channels instead

Suggested fix

pi0_pytorch.py is the only caller of preprocess_observation_pytorch, and it feeds SigLIP, which requires [B, C, H, W]. Making that the unconditional output pins the contract in one place:

-        # Convert back to [B, C, H, W] format if it was originally channels-first
-        if is_channels_first:
-            image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]
+        # SigLIP consumes [B, C, H, W], so always emit that regardless of the input layout.
+        image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]

Fixing it in model.py instead would break the wandb image logging at scripts/train_pytorch.py#L378, which assumes Observation.images is channels-first. That logging path is already broken for NHWC inputs and is tracked separately in #877.

Bug 2: vision_config.projection_dim is hardcoded to 2048

With bug 1 fixed, the same command fails again. In PaliGemmaWithExpertModel.__init__, every text-side field is derived from the gemma variant, but the vision projection width is a literal:

https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/models_pytorch/gemma_pytorch.py#L26-L38

vlm_config_hf.text_config.hidden_size = vlm_config.width
...
vlm_config_hf.vision_config.projection_dim = 2048

2048 is the width of gemma_2b, so the coincidence hides the bug for every shipped training config. A variant with any other width produces a projector that writes into the wrong embedding width:

dummy:      text hidden_size=64,   vision projection_dim=2048
gemma_300m: text hidden_size=1024, vision projection_dim=2048

Scope

This is narrower than bug 1. Every paligemma_variant in src/openpi/training/config.py is either gemma_2b or gemma_2b_lora, both width 2048, so no shipped training config is affected. What breaks today is the debug, debug_restore, and debug_pi05 configs, which use paligemma_variant="dummy" (width 64), plus any custom config choosing a non-2b variant. Since paligemma_variant is a public field on Pi0Config, the constraint is currently implicit.

The JAX path derives this value rather than hardcoding it:

https://github.com/Physical-Intelligence/openpi/blob/215abfb/src/openpi/models/pi0.py#L82-L83

_siglip.Module(
    num_classes=paligemma_config.width,

Traceback

  File "src/openpi/models_pytorch/pi0_pytorch.py", line 331, in forward
    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)
  File "src/openpi/models_pytorch/pi0_pytorch.py", line 228, in embed_prefix
    embs = torch.cat(embs, dim=1)
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 2048 but got size 64 for tensor number 3 in the list.

Reproduction without the training script

from openpi.models import gemma
from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel

for variant in ["dummy", "gemma_300m"]:
    config = gemma.get_config(variant)
    model = PaliGemmaWithExpertModel(config, gemma.get_config("dummy"), precision="float32")
    print(variant, model.paligemma.config.text_config.hidden_size, model.paligemma.config.vision_config.projection_dim)

Suggested fix

-        vlm_config_hf.vision_config.projection_dim = 2048
+        vlm_config_hf.vision_config.projection_dim = vlm_config.width

This is a no-op for gemma_2b and gemma_2b_lora, so released checkpoints are unaffected.

Verification

Applying both fixes lets the documented command finish:

Training: 100%|██████████| 10/10 [00:23<00:00, 2.31s/it, loss=3.0191, lr=2.50e-07, step=10]
Saved checkpoint at step 9 -> checkpoints/debug/<exp_name>/9

I confirmed each bug in isolation on clean 215abfb: reverting only the layout fix reproduces the conv2d error, and reverting only the projection_dim fix reproduces the torch.cat error.

The layout fix does not disturb the uint8 path that real datasets use. Both dtypes converge on the same shape and range:

uint8 NHWC (real dataset)      from_dict=(2, 3, 224, 224) -> preprocess=(2, 3, 224, 224) range=[-1.00, 0.99]
float32 NHWC (FakeDataset)     from_dict=(2, 224, 224, 3) -> preprocess=(2, 3, 224, 224) range=[-1.00, 1.00]

Why this was not caught

The PyTorch path has no test coverage. scripts/train_test.py only exercises the JAX train.py, and no test imports train_pytorch, pi0_pytorch, or models_pytorch:

$ git grep -ln "train_pytorch\|pi0_pytorch\|models_pytorch" 215abfb -- '*_test.py'
(no matches)

Offer

I have both fixes plus a regression test locally (66 passed, up from 59, with uv run pytest --strict-markers -m "not manual"). Two of the new tests need transformers_replace in site-packages, so they skip on a plain uv sync environment unless CI adds that copy step. Happy to open a PR if useful.