Update modeling_tri_stream.py

modeling_tri_stream.py

@@ -182,14 +182,14 @@ def procrustes_align(source, target, whiten=False):
         target: (N, D) — target positions (e.g. class centroids)
         whiten: if True, normalize variance per dimension before alignment
     """
-    source_c = source - source.mean(0, keepdim=True)
-    target_c = target - target.mean(0, keepdim=True)
+    source_c = source.float() - source.float().mean(0, keepdim=True)
+    target_c = target.float() - target.float().mean(0, keepdim=True)
 
     if whiten:
         source_c = source_c / (source_c.std(0, keepdim=True) + 1e-8)
         target_c = target_c / (target_c.std(0, keepdim=True) + 1e-8)
 
-    M = source_c.T @ target_c  # (D, D)
+    M = (source_c.T @ target_c).float()  # SVD requires float32
     U, S, Vt = torch.linalg.svd(M)
 
     # Ensure proper rotation (det = +1)
@@ -658,33 +658,30 @@ class TriStreamViT(nn.Module):
         3. Procrustes-align matched anchors to centroids
         4. Apply rotation to ALL anchors as small step
         """
-        centroids = simplex_buffer.class_centroids(self.num_classes)
-        if centroids is None:
-            return None
+        # Force float32 — SVD and det don't support bf16
+        with torch.amp.autocast("cuda", enabled=False):
+            centroids = simplex_buffer.class_centroids(self.num_classes)
+            if centroids is None:
+                return None
 
-        anchors = self.gal.gal_anchors  # (A, D)
+            anchors = self.gal.gal_anchors.float()  # (A, D)
 
-        # Match: for each centroid, find nearest anchor
-        centroid_n = F.normalize(centroids, dim=-1)
-        anchor_n = F.normalize(anchors, dim=-1)
-        cos = centroid_n @ anchor_n.T  # (C, A)
-        matched_idx = cos.argmax(dim=1)  # (C,) — one anchor per class
+            centroid_n = F.normalize(centroids.float(), dim=-1)
+            anchor_n = F.normalize(anchors, dim=-1)
+            cos = centroid_n @ anchor_n.T
+            matched_idx = cos.argmax(dim=1)
 
-        # Extract matched anchor positions
-        matched_anchors = anchors[matched_idx]  # (C, D)
+            matched_anchors = anchors[matched_idx]
 
-        # Procrustes: rotate matched_anchors → centroids
-        R, score = procrustes_align(
-            matched_anchors.float(), centroids.float(), whiten=whiten)
-        R = R.to(anchors.dtype)
+            R, score = procrustes_align(
+                matched_anchors, centroids.float(), whiten=whiten)
 
-        # Apply as small rotation step to ALL anchors
-        rotated = anchors @ R
-        self.gal.gal_anchors.copy_(
-            F.normalize(anchors + lr * (rotated - anchors), dim=-1)
-        )
+            rotated = anchors @ R
+            new_anchors = F.normalize(
+                anchors + lr * (rotated - anchors), dim=-1)
+            self.gal.gal_anchors.copy_(new_anchors.to(self.gal.gal_anchors.dtype))
 
-        return score
+            return score
 
     # ──────────────────────────────────────────────────────────
     # LOSS COMPUTATION