From 8197ab24072501a843f70108891aaf1284180617 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 16:38:16 +0300 Subject: [PATCH] Fix training loop: only pass positive_weights to WeightedInfoNCELoss MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/training/train_gtauav.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 89048c2..4e1b16e 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -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: