diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py
index 7a85190b71b1..f4ec14bf4990 100644
--- a/deepspeed/ops/transformer/inference/ds_attention.py
+++ b/deepspeed/ops/transformer/inference/ds_attention.py
@@ -244,8 +244,7 @@ def compute_attention(self, qkv_out, input_mask, layer_past, alibi):
         ) * self.num_attention_heads_per_partition if dist.is_initialized() else 0
         attention_probs = self.softmax_func(
             attn_scores=attention_scores,
-            attn_mask=((1 - input_mask).half() *
-                       minus_inf) if input_mask.dtype == torch.int64 else input_mask,
+            attn_mask=((1 - input_mask).half() * minus_inf),
             alibi=alibi,
             triangular=(self.config.triangular_masking
                         and (attention_scores.shape[-2] > 1)),