bugfix

app/models/heads/mask2former_vid.py

@@ -190,7 +190,7 @@ class Mask2FormerVideoHead(AnchorFreeHead):
                 #     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 = torch.zeros(1, _dim, dtype=torch.float32, device='cpu')
                 # back_token = back_token.to(device='cpu')
                 cls_embed = torch.cat([
                     cls_embed, back_token.repeat(_prototypes, 1)[None]
@@ -597,7 +597,7 @@ class Mask2FormerVideoHead(AnchorFreeHead):
             input_query_bbox = input_bbox_embed
 
         tgt_size = pad_size
-        attn_mask = torch.ones(tgt_size, tgt_size).to('cuda') < 0
+        attn_mask = torch.ones(tgt_size, tgt_size).to(input_bbox_embed.device) < 0
         # match query cannot see the reconstruct
         attn_mask[pad_size:, :pad_size] = True
         # reconstruct cannot see each other

app/models/heads/yoso_head.py

@@ -376,7 +376,7 @@ class CrossAttenHead(nn.Module):
             #     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 = torch.zeros(1, _dim, dtype=torch.float32, device='cpu')
             # back_token = back_token.to(device='cpu')
             cls_embed = torch.cat([
                 cls_embed, back_token.repeat(_prototypes, 1)[None]