Fix training loop: only pass positive_weights to WeightedInfoNCELoss

InfoNCELoss.forward() doesn't accept `positive_weights` — that kwarg is
specific to WeightedInfoNCELoss's adaptive label-smoothing path. After
switching the default `loss_type` to "symmetric", training crashed with
`TypeError: unexpected keyword argument 'positive_weights'`.

Build the kwargs dict conditionally: add `positive_weights` only when
the loss is an instance of WeightedInfoNCELoss.

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

View File

@@ -825,16 +825,19 @@ def train(cfg: TrainConfigGTAUAV) -> None:
sat_caption_l2=batch["sat_caption_l2"],
sat_caption_l3=batch["sat_caption_l3"],
)
# Loss — WeightedInfoNCE with positive weights from dataset.
pos_weights = batch["positive_weights"].to(cfg.device, non_blocking=True)
# Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
queue_neg = neg_bank.get_queue() if neg_bank is not None else None
loss_dict = loss_fn(
embeddings=embeddings,
epoch=epoch,
total_epochs=cfg.epochs,
positive_weights=pos_weights,
queue_negatives=queue_neg,
)
loss_kwargs = {
"embeddings": embeddings,
"epoch": epoch,
"total_epochs": cfg.epochs,
"queue_negatives": queue_neg,
}
if isinstance(loss_fn, WeightedInfoNCELoss):
loss_kwargs["positive_weights"] = batch["positive_weights"].to(
cfg.device, non_blocking=True,
)
loss_dict = loss_fn(**loss_kwargs)
# Enqueue current gallery embeddings (detached).
if neg_bank is not None: