fix: minor error

jina-ai · Aug 10, 2022 · a72b3fd · a72b3fd
1 parent e30aeb0
commit a72b3fd
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/server/clip_server/model/model.py b/server/clip_server/model/model.py
@@ -294,7 +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):
-        attn_mask = attn_mask.to(dtype=x.dtype, device=x.device) if attn_mask else 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