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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user