diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
index 78cbcce2c..e1c5e5164 100644
--- 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