fix: cast dtype for fp16 (#801)

* fix: cast dtype for fp16 * fix: minor error
jina-ai · Aug 10, 2022 · bb2c142 · bb2c142
1 parent 25c91e2
commit bb2c142
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -294,6 +294,8 @@ 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):
+        if attn_mask is not None:
+            attn_mask = attn_mask.to(dtype=x.dtype, device=x.device)
         x = x + self.attention(self.ln_1(x), attn_mask=attn_mask)
         x = x + self.mlp(self.ln_2(x))
         return x