fix: ScaledFP8WeightTensor crashes on detach/to when inner tensors are inference tensors

[Gunther-Schulz]

I created this PR to be able to run a fine-tuned LTX2.3 Dev model that's only available as FP8.

🤖 Generated with Claude Code

Problem

Running FP8-quantized checkpoints (e.g. LTX-2.3 22B) crashes with:

RuntimeError: Cannot set version_counter for inference tensor

The crash occurs in ScaledFP8WeightTensor.__torch_function__ when LTX-2's _duplicate_timestep_embedder calls .detach().to(device) on a quantized weight during the prewarm step. A secondary crash follows in _linear_fallback:

RuntimeError: Expected all tensors to be on the same device

Root cause

ScaledFP8WeightTensor is a _make_wrapper_subclass tensor whose inner _data (fp8) and _scale (float32) tensors are created as inference tensors during model loading (under torch.inference_mode()). The wrapper itself is not an inference tensor, but its inner tensors are.

When .detach() is called on the wrapper, PyTorch attempts to set up version-counter sharing on the inner tensors — which is illegal for inference tensors. Two dispatch layers were affected:

1. __torch_function__: receives the C-level TensorBase.detach method (not torch.Tensor.detach), so identity checks fail. Must match by func.__name__.
2. __torch_dispatch__: the existing aten.detach handler called op(t._data) directly, which triggers the same version-counter error.
3. __torch_dispatch__ _to_copy: same issue — called op(t._data) without guarding against inference tensors.

The _linear_fallback device mismatch was a separate latent bug: the fallback path converted weights to the right dtype but forgot to move them to the right device, only surfacing when the fallback was actually exercised.

Fix

• __torch_function__: intercept detach and to by name before they reach C++ dispatch. For detach, return a new wrapper sharing the same inner tensors (safe — model weights never require grad). For to, use torch.inference_mode(False) to move inner tensors without version-counter errors.
• __torch_dispatch__ detach: stop calling op(t._data) / op(t._scale); share inner tensors directly instead.
• __torch_dispatch__ _to_copy: wrap op(t._data) / op(t._scale) in torch.inference_mode(False) for consistency with the __torch_function__ path.
• _linear_fallback: capture target_device = input.device and pass it to both .to() calls so weights land on the correct device.

Why this only surfaces with FP8 checkpoints

The bug requires ScaledFP8WeightTensor to exist, which only happens for FP8-quantized weights. bf16 checkpoints never create these wrappers so the crash path is never reached.

Testing

Verified end-to-end generation with LTX-2.3 22B FP8 checkpoint. No regressions expected for bf16 models as the changed code paths are only reached for ScaledFP8WeightTensor instances.

🤖 Generated with Claude Code