ByteDance-Seed / ByteDance-Seed/Depth-Anything-3
prepare_cls_token fails on 2nd+ inference when called under nested torch.inference_mode() (nn.Parameter cls_token)
- Dominant language
- Python
- Stars
- 6.3k
- Forks
- 702
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
When `DepthAnything3.forward` (decorated with `@torch.inference_mode()`) is called **more than once** inside a process that itself is already under `torch.inference_mode()` — for example inside a ComfyUI persistent worker that wraps every model call in `torch.inference_mode()` — the **second call fails** with:
```
RuntimeError: Inference tensors do not track version counter.
File ".../depth_anything_3/model.py", line 668, in prepare_cls_token
cls_token = self.cls_token.expand(B, S, -1)
```
## Reproduction
1. Load `DepthAnything3.from_pretrained(...)` (any preset)
2. Call `model.inference([img1])` — succeeds
3. Call `model.inference([img2])` — fails at `prepare_cls_token` with the above error
If the call is wrapped in an additional `torch.inference_mode()` context (as comfy_env's persistent worker does — see comfy_env#6), this happens **100% of the time** on the 2nd+ call.
## Root cause
`model.py:583`:
```python
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
```
`@torch.inference_mode()` on `forward()` (api.py:99) converts `cls_token` into an "inference tensor" on first call. Inference tensors **do not track version counters**. On the second call, downstream `.expand()` / `.to()` ops need to check the version counter (autograd metadata) and fail.
This is the documented PyTorch #90882 pattern: **`torch.inference_mode()` is not safe to compose with itself**, and it leaves persistent state on `Parameter` buffers.
## Suggested fix
Either:
1. **Replace `@torch.inference_mode()` with `@torch.no_grad()` on the public `forward()` (line 99).** `no_grad` disables autograd without the strict "inference mode" treatment that strips version counters. This is the canonical fix for PyTorch #90882. Marginal perf cost (~10-15%) but safe.
2. **Don't decorate `forward()` at all** — let the caller choose their own gradient context. The `@torch.no_grad()` block at line 127 (inside `inference()`) is sufficient to disable autograd for the public API.
3. **Re-create `cls_token` (or any "inference-stamped" buffer) per call** inside the inference wrapper — heavy-handed.
Option 1 is what PyTorch upstream recommends for any user-facing inference decorator.
## Verified locally
Patching `api.py:99` from `@torch.inference_mode()` to `@torch.no_grad()` resolves the issue completely. With the patch, repeated `model.inference()` calls inside an `inference_mode()`-wrapped context (comfy_env persistent worker) all succeed.
## Environment
- DA3 commit: `main` (latest)
- PyTorch 2.10.0+cu128
- comfy-env v0.2.x persistent worker (comfy_env#6)
Thanks for the great model — this is otherwise a beautiful piece of code, and the fix is one line.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start at api.py:99, where DepthAnything3.forward is decorated, then inspect model.py:583 and prepare_cls_token at line 668. Reproduce repeated model.inference calls inside nested torch.inference_mode() using the reported environment. Done means subsequent calls succeed without the inference-tensor version-counter error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100