mypy fixes

doctr/models/modules/transformer/pytorch.py

@@ -46,7 +46,7 @@ def scaled_dot_product_attention(
     query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Scaled Dot-Product Attention"""
-    scores: torch.Tensor = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))
+    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))  # type: ignore[assignment]
     if mask is not None:
         # NOTE: to ensure the ONNX compatibility, masked_fill works only with int equal condition
         scores = scores.masked_fill(mask == 0, float("-inf"))

doctr/transforms/functional/pytorch.py

@@ -30,7 +30,7 @@ def invert_colors(img: torch.Tensor, min_val: float = 0.6) -> torch.Tensor:
     out = F.rgb_to_grayscale(img, num_output_channels=3)
     # Random RGB shift
     shift_shape = [img.shape[0], 3, 1, 1] if img.ndim == 4 else [3, 1, 1]
-    rgb_shift: torch.Tensor = min_val + (1 - min_val) * torch.rand(shift_shape)
+    rgb_shift = min_val + (1 - min_val) * torch.rand(shift_shape)  # type: ignore[assignment]
     # Inverse the color
     if out.dtype == torch.uint8:
         out = (out.to(dtype=rgb_shift.dtype) * rgb_shift).to(dtype=torch.uint8)

doctr/transforms/modules/pytorch.py

@@ -193,8 +193,8 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor:
         # Reshape the distribution
         try:
             if x.dtype == torch.uint8:
-                return (  # type: ignore[attr-defined]
-                    (
+                return (
+                    (  # type: ignore[attr-defined]
                         255
                         * random_shadow(
                             x.to(dtype=torch.float32) / 255,