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: