Fix resize_with_pad_torch dropping the batch dimension for a single-image batch

[ousamabenyounes]

Summary

Fix #805.

resize_with_pad_torch dropped the batch dimension for a single-image batch,
diverging from the JAX resize_with_pad, which preserves the input rank.

For a channels-last 4D input of shape [1, H, W, C] the function returned a 3D
[H', W', C] tensor, because the final block squeezed dim 0 whenever
batch_size == 1 — regardless of whether the batch dimension was genuinely
present in the input or had been added internally for a 3D image:

if channels_last:
    padded_images = padded_images.permute(0, 2, 3, 1)
    if batch_size == 1 and images.shape[0] == 1:   # always true for a size-1 batch
        padded_images = padded_images.squeeze(0)

This breaks callers that feed batched observations (e.g. Observation.from_dict
with an image of shape (1, 3, 480, 640)) into the PyTorch path, and is
inconsistent with the JAX reference which keeps [1, H, W, C] as [1, H', W', C].

Fix

Track whether a batch dimension was added internally (mirroring the JAX
has_batch_dim) and squeeze it back only in that case, for both channels-last
and channels-first inputs. A genuine size-1 batch is now preserved.

Test verification (RED → GREEN)

New regression test test_resize_with_pad_torch_preserves_batch_dim in
src/openpi/shared/image_tools_test.py.

RED — new test on main with the fix reverted (prod code only):

>       assert tuple(resized.shape) == (1, 224, 224, 3)
E       assert (224, 224, 3) == (1, 224, 224, 3)
src/openpi/shared/image_tools_test.py:47: AssertionError
============================== 1 failed in 2.14s ===============================

GREEN — with the fix:

src/openpi/shared/image_tools_test.py::test_resize_with_pad_shapes PASSED [ 50%]
src/openpi/shared/image_tools_test.py::test_resize_with_pad_torch_preserves_batch_dim PASSED [100%]
============================== 2 passed in 3.46s ===============================

ruff check . and ruff format --check are clean on the changed files.

Files changed

File	Change
src/openpi/shared/image_tools.py	Track added batch dim; squeeze back only when added
src/openpi/shared/image_tools_test.py	Regression test for batch-dim preservation

Full local suite

Local CI mirrors .github/workflows/{test,pre-commit}.yml
(ruff check ., ruff format --check, uv run pytest --strict-markers -m "not manual"):

• ruff check . — all checks passed.
• ruff format --check — clean on the changed files (2 pre-existing, unrelated
  drifted files in examples/ are unchanged here).
• Tests: 55 passed across every suite that fits in local RAM, including the
  new regression test. Two GPU-scale model-build tests
  (src/openpi/models/model_test.py, scripts/train_test.py::test_train[debug])
  are OOM-killed on the local 15 GiB host — identically on the unmodified main
  baseline — so this is a pre-existing environment limit, not a regression from
  this change. This diff only touches image_tools, which those tests do not
  exercise; GitHub CI runs the full matrix.