fix: cast dtype for fp16

jina-ai · Aug 10, 2022 · e30aeb0 · e30aeb0
1 parent 25c91e2
commit e30aeb0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -294,6 +294,7 @@ def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
         return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
 
     def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        attn_mask = attn_mask.to(dtype=x.dtype, device=x.device) if attn_mask else None
         x = x + self.attention(self.ln_1(x), attn_mask=attn_mask)
         x = x + self.mlp(self.ln_2(x))
         return x