Fix autograd in-place error: move memory-bank enqueue after backward

Hard mining (c307269) uses `queue.t()` as a view directly in the
similarity matmul, so the autograd graph now holds a reference to the
memory-bank buffer's storage. The previous code enqueued fresh gallery
embeddings BEFORE `backward()`, which mutates that same buffer in place
and triggers:

  RuntimeError: one of the variables needed for gradient computation has
  been modified by an inplace operation: [torch.cuda.FloatTensor [1024, 8]]
  is at version 2; expected version 1 instead.

The older InfoNCE path called `torch.cat([emb_b, queue], 0)`, which
materialises a fresh contiguous tensor and severs the alias — so the
same enqueue order worked. After the mining refactor, we need the
buffer to stay stable until backward completes.

Moving the enqueue to after backward is semantically identical: the
queue state observed by the next training step is the same either way
(FIFO, not involved in the just-completed backward).

`smoke_train.py` already had the correct order and didn't catch the
regression.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 16:40:58 +03:00
parent 8197ab2407
commit 70f1617317

View File

@@ -839,15 +839,20 @@ def train(cfg: TrainConfigGTAUAV) -> None:
)
loss_dict = loss_fn(**loss_kwargs)
# Enqueue current gallery embeddings (detached).
if neg_bank is not None:
neg_bank.enqueue(embeddings["gallery"].detach())
# Scale loss by accumulation steps so gradients average correctly.
raw_loss = float(loss_dict["total"].item()) # save before backward
total_loss = loss_dict["total"] / accum
scaler.scale(total_loss).backward()
# Enqueue current gallery AFTER backward. The queue buffer is aliased
# into the autograd graph through `queue_neg` (a view returned by
# `NegativeMemoryBank.get_queue`), so modifying it before backward
# triggers "variable needed for gradient computation has been modified
# by an inplace operation". Enqueueing here is semantically identical
# — the next step's queue state is the same either way.
if neg_bank is not None:
neg_bank.enqueue(embeddings["gallery"].detach())
# Optimizer step only after accumulating `accum` micro-batches.
is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader)
if is_accum_step: