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)

Open Beginner friendly
#266 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.