code review, fixing a small mypy issue

facebookresearch · Nov 4, 2021 · 72defc3 · 72defc3
1 parent 013a927
commit 72defc3
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/xformers/components/attention/nystrom.py b/xformers/components/attention/nystrom.py
@@ -180,8 +180,6 @@ def forward(
                 )
                 key_padding_mask = bool_mask_to_additive(key_padding_mask)
 
-            assert key_padding_mask is not None  # mypy is drunk
-
             if key_padding_mask.ndim == 2:
                 key_padding_mask = reshape_key_padding_mask(
                     key_padding_mask, batched_dim

diff --git a/xformers/components/attention/utils.py b/xformers/components/attention/utils.py
@@ -91,7 +91,7 @@ def iterative_pinv(softmax_mat: torch.Tensor, n_iter=6, pinverse_original_init=F
 
 def bool_mask_to_additive(
     mask: torch.Tensor, dtype: Optional[torch.dtype] = torch.float32
-):
+) -> torch.Tensor:
     assert (
         mask.dtype == torch.bool
     ), "This util is meant to convert in between bool masks and additive ones"