don't use subtraction of boolean tensors, use xor (#420)

* don't use subtraction of boolean tensors, use xor This code fails to run without the change * rustfmt --------- Co-authored-by: Charles Samuels <[email protected]> Co-authored-by: guillaume-be <[email protected]>
guillaume-be · Sep 30, 2023 · 98c7905 · 98c7905
1 parent 9575902
commit 98c7905
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/models/bert/bert_model.rs b/src/models/bert/bert_model.rs
@@ -396,9 +396,11 @@ impl<T: BertEmbedding> BertModel<T> {
             train,
         )?;
 
-        let extended_attention_mask: Tensor =
-            ((extended_attention_mask.ones_like() - extended_attention_mask) * -10000.0)
-                .to_kind(embedding_output.kind());
+        let extended_attention_mask: Tensor = ((extended_attention_mask
+            .ones_like()
+            .bitwise_xor_tensor(&extended_attention_mask))
+            * -10000.0)
+            .to_kind(embedding_output.kind());
 
         let encoder_extended_attention_mask: Option<Tensor> =
             if self.is_decoder & encoder_hidden_states.is_some() {