bugfix

app/models/heads/mask2former_vid.py

@@ -183,14 +183,15 @@ class Mask2FormerVideoHead(AnchorFreeHead):
                 _dim = cls_embed.size(2)
                 _prototypes = cls_embed.size(1)
 
-                if rank == 0:
-                    back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
-                    # back_token = back_token / back_token.norm(p=2, dim=-1, keepdim=True)
-                else:
-                    back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
-                if world_size > 1:
-                    dist.broadcast(back_token, src=0)
-                back_token = back_token.to(device='cpu')
+                # if rank == 0:
+                #     back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+                #     # back_token = back_token / back_token.norm(p=2, dim=-1, keepdim=True)
+                # else:
+                #     back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
+                # if world_size > 1:
+                #     dist.broadcast(back_token, src=0)
+                back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+                # back_token = back_token.to(device='cpu')
                 cls_embed = torch.cat([
                     cls_embed, back_token.repeat(_prototypes, 1)[None]
                 ], dim=0)

app/models/heads/yoso_head.py

@@ -369,13 +369,15 @@ class CrossAttenHead(nn.Module):
             # background class
             _dim = cls_embed.size(2)
             _prototypes = cls_embed.size(1)
-            if rank == 0:
-                back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
-            else:
-                back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
-            if world_size > 1:
-                dist.broadcast(back_token, src=0)
-            back_token = back_token.to(device='cpu')
+            # if rank == 0:
+            #     back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+            #     # back_token = back_token / back_token.norm(p=2, dim=-1, keepdim=True)
+            # else:
+            #     back_token = torch.empty(1, _dim, dtype=torch.float32, device='cuda')
+            # if world_size > 1:
+            #     dist.broadcast(back_token, src=0)
+            back_token = torch.zeros(1, _dim, dtype=torch.float32, device='cuda')
+            # back_token = back_token.to(device='cpu')
             cls_embed = torch.cat([
                 cls_embed, back_token.repeat(_prototypes, 1)[None]
             ], dim=0)