Fix NaN: compute loss in fp32 outside AMP autocast

Root cause: GradScaler scales gradients by ~65536 in fp16, causing
logit_scale.exp() gradient to overflow. The learnable temperature
and similarity logits must stay in fp32.

Fix: model forward runs inside autocast(fp16), but loss computation
(similarity @ temperature + cross_entropy) runs outside in fp32.

Also: clamp logit_scale in logit-space before exp() and force
similarity computation to fp32.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/training/train_gtauav.py

@@ -361,6 +361,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
             sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
 
+            # Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
             with autocast(device_type="cuda", enabled=cfg.use_amp):
                 if cfg.baseline_mode:
                     embeddings = model(drone_img=drone_img, sat_img=sat_img)
@@ -372,11 +373,12 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                         caption_l2=batch["caption_l2"],
                         caption_l3=batch["caption_l3"],
                     )
-                loss_dict = loss_fn(
-                    embeddings=embeddings,
-                    epoch=epoch,
-                    total_epochs=cfg.epochs,
-                )
+            # Loss in fp32 (learnable temperature gradient overflows in fp16).
+            loss_dict = loss_fn(
+                embeddings=embeddings,
+                epoch=epoch,
+                total_epochs=cfg.epochs,
+            )
 
             total_loss = loss_dict["total"]
             scaler.scale(total_loss).backward()