update

demo.py

@@ -227,7 +227,7 @@ def visual_grounding(image, text):
     sample = utils.move_to_cuda(sample) if use_cuda else sample
     sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
 
-    with torch.no_grad():
+    with torch.autocast("cpu"), torch.no_grad():
         if isinstance(models, list):
             model = models[0]
         min_len = 6

models/polyformer/swin.py

@@ -287,7 +287,7 @@ class WindowAttention(nn.Module):
         else:
             attn = self.softmax(attn)
 
-        attn = self.attn_drop(attn).half()
+        attn = self.attn_drop(attn)
 
         x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
         x = self.proj(x)
@@ -758,4 +758,4 @@ class SwinTransformer(nn.Module):
         #outs = outs['layer2'] # [B, 512, 32, 32]
         outs = outs[f'layer{self.out_indices[0]}']  # [B, 512, 16, 16]
         #print(outs.shape)
-        return outs
+        return outs