diff --git a/layers/common_layers.py b/layers/common_layers.py
index 2edf0dab..4269728f 100644
--- a/layers/common_layers.py
+++ b/layers/common_layers.py
@@ -240,7 +240,7 @@ class Attention(nn.Module):
                 attention_hidden_state, processed_inputs)
         # apply masking
         if mask is not None:
-            attention.data.masked_fill_(1 - mask, self._mask_value)
+            attention.data.masked_fill_(torch.bitwise_not(mask), self._mask_value)
         # apply windowing - only in eval mode
         if not self.training and self.windowing:
             attention = self.apply_windowing(attention, inputs)