fix(models): build spatial_shapes from Python ints under torch.compile - #1411

What this fixes

compile=True aborts the training step. Transformer.forward builds spatial_shapes with

spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(
    device=srcs[0].device, dtype=torch.long
)

and Dynamo polyfills torch._shape_as_tensor to return a torch.Size instead of a tensor, so torch.stack raises:

TypeError: expected Tensor as element 0 in argument 0, but got torch.Size
  src/rfdetr/models/transformer.py:346 in forward

Minimal reproduction, no rf-detr involved:

import torch

def f(x):
    return torch.stack([torch._shape_as_tensor(x)[2:4]]).to(device=x.device, dtype=torch.long)

x = torch.randn(1, 3, 8, 8, device="cuda")
f(x)                               # eager: tensor([[8, 8]])
torch.compile(f, dynamic=True)(x)  # TypeError: ... but got torch.Size

And through the public API:

RFDETRNano(compile=True).train(dataset_dir=..., multi_scale=False)
# TypeError: expected Tensor as element 0 in argument 0, but got torch.Size

torch._dynamo.config.suppress_errors = True (set at module_model.py:436) does not absorb it. That flag only swallows errors Dynamo itself raises; here the polyfill makes user code raise a real TypeError, so Dynamo propagates it. Measured both ways, True and False, same result.

Reproduced on torch 2.9.1+cu129 (L4) and 2.13.0+cu130 (RTX 4060 Laptop), and at dynamic=True, dynamic=None and dynamic=False alike, so this is not a recent regression and not a dynamic-shapes problem.

Under the shipped default multi_scale=True the gate at module_model.py:417-418 disables compilation before this line is reached, so a default compile=True run does not crash .. it just runs eager and logs one line. The crash needs multi_scale=False.

The fix

The branch that works already exists two lines above. torch.export/ExecuTorch cannot trace torch._shape_as_tensor either, and #1142 handled that by building spatial_shapes from the Python-int (H, W) pairs the function already collects at transformer.py:292-295. torch.compile needs the same branch for a different reason, so the condition now covers both:

if (
    getattr(torch.compiler, "is_exporting", _tracer_absent)()
    or getattr(torch.compiler, "is_compiling", _tracer_absent)()
):
    spatial_shapes = torch.as_tensor(spatial_shapes_hw, device=srcs[0].device, dtype=torch.long)
else:
    spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(...)

The _shape_as_tensor form stays for every path that needs it. #1155 introduced it because TensorRT accepts the Constant it produces while a ScatterND-producing form is rejected, and that path goes through torch.jit.trace, where both guards are false. Checked directly rather than assumed:

getattr guards the same way the existing is_exporting call does, for torch versions that lack the attribute. The existing private fallback was named _not_exporting, which would read wrong on the second call, so it is now _tracer_absent with a docstring covering both predicates. It has no other caller and is not exported.

torch.compiler.is_compiling() was added in torch 2.3 and is_exporting() in torch 2.7 (checked against the upstream source for each release). Below torch 2.3 the getattr fallback returns False, so this branch cannot detect torch.compile there and compile=True keeps raising the same TypeError develop already raises — unchanged by this fix. There is no working substitute on torch<2.3: both torch._dynamo.is_compiling() and torch._utils.is_compiling() are hardcoded return False stubs on torch 2.2.x rather than real predicates, so torch 2.2.x (this project's declared floor) is unaffected either way, the same as it was already unaffected by is_exporting().

The same public call that raises on develop now runs to the end:

RFDETRNano(num_classes=80, pretrain_weights=None, compile=True).train(
    dataset_dir=..., epochs=1, batch_size=2, num_workers=0, multi_scale=False,
)
# develop: TypeError ... but got torch.Size
# with this change: train() completes

Tests

Two in tests/models/test_transformer_onnx_spatial_shapes.py, the file that already owns this line:

• test_spatial_shapes_survives_dynamo_shape_as_tensor_polyfill reproduces the polyfill (torch._shape_as_tensor patched to return a torch.Size) together with is_compiling() == True, runs the real Transformer.forward and checks the output still matches the eager one. Parametrised over one and two feature levels, since that is what changes the shape of the source spatial_shapes is built from. Both cases fail on develop with the exact production error and pass with the fix. Neither invokes the compiler, so they cost nothing in CI.
• test_spatial_shapes_compile_guard_is_false_under_torchscript_trace pins the guard that keeps the TensorRT path on the old formulation.

The file's module-level onnx import moved to a per-fixture pytest.importorskip inside the three fixtures that build an ONNX graph, so these two tests (and the pre-existing test_level_start_index_correctness_two_levels) collect and run under ci-tests-cpu.yml's train,augment,cli,visual install set, which does not include the [onnx] extra.

pytest tests/models/test_transformer_onnx_spatial_shapes.py is 10 passed, and the whole tests/models with the CI markers is 935 passed / 6 skipped. pytest tests/export tests/training/test_module_model.py with the same markers is 623 passed / 48 skipped / 1 failed, the failure being tests/export/test_onnx_notes.py::_export_tiny_model, a pre-existing doctest failure on develop unrelated to this change. pre-commit run --all-files clean, including mypy --strict.

Scope

This makes the compile path reachable. It does not claim a speedup: no throughput number can be produced until compilation runs at all, and the measurement belongs in #1410. What I can say from that side is that the workload is worth compiling, the training step issues ~4200-4800 CUDA kernels regardless of batch size, and on a g2-standard-8 L4 the GPU is busy only 19-46% of the wall time.