Format bert.py

huggingface · Oct 20, 2023 · d7c3fde · d7c3fde
1 parent 3f21ebb
commit d7c3fde
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/candle-pyo3/py_src/candle/models/bert.py b/candle-pyo3/py_src/candle/models/bert.py
@@ -59,8 +59,7 @@ def forward(self, hidden_states: Tensor, attention_mask=None) -> Tensor:
         attention_scores = attention_scores / float(self.attention_head_size) ** 0.5
         if attention_mask is not None:
             b_size, _, _, last_dim = attention_scores.shape
-            attention_scores = attention_scores.broadcast_add(
-                    attention_mask.reshape((b_size, 1, 1, last_dim)))
+            attention_scores = attention_scores.broadcast_add(attention_mask.reshape((b_size, 1, 1, last_dim)))
         attention_probs = F.softmax(attention_scores, dim=-1)
 
         context_layer = attention_probs.matmul(value)
@@ -198,7 +197,9 @@ def __init__(self, config: Config, add_pooling_layer=True) -> None:
         self.encoder = BertEncoder(config)
         self.pooler = BertPooler(config) if add_pooling_layer else None
 
-    def forward(self, input_ids: Tensor, token_type_ids: Tensor, attention_mask=None) -> Tuple[Tensor, Optional[Tensor]]:
+    def forward(
+        self, input_ids: Tensor, token_type_ids: Tensor, attention_mask=None
+    ) -> Tuple[Tensor, Optional[Tensor]]:
         if attention_mask is not None:
             # Replace 0s with -inf, and 1s with 0s.
             attention_mask = masked_fill(float("-inf"), attention_mask, 1.0)