Fix docs and mypy

Author: datumbox

torchvision/models/video/mvit.py

@@ -290,9 +290,9 @@ def __init__(
                 norm_layer(self.head_dim),
             )
 
-        self.rel_pos_h: Optional[nn.Module] = None
-        self.rel_pos_w: Optional[nn.Module] = None
-        self.rel_pos_t: Optional[nn.Module] = None
+        self.rel_pos_h: Optional[nn.Parameter] = None
+        self.rel_pos_w: Optional[nn.Parameter] = None
+        self.rel_pos_t: Optional[nn.Parameter] = None
         if rel_pos:
             assert input_size[1] == input_size[2]  # TODO: remove this limitation
             size = input_size[1]
@@ -471,6 +471,8 @@ def __init__(
             temporal_size (int): The temporal size ``T`` of the input.
             block_setting (sequence of MSBlockConfig): The Network structure.
             residual_pool (bool): If True, use MViTv2 pooling residual connection.
+            rel_pos (bool): TODO
+            dim_mul_in_att (bool): TODO
             dropout (float): Dropout rate. Default: 0.0.
             attention_dropout (float): Attention dropout rate. Default: 0.0.
             stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0.
@@ -508,7 +510,7 @@ def __init__(
         # Spatio-Temporal Class Positional Encoding
         self.pos_encoding = PositionalEncoding(
             embed_size=block_setting[0].input_channels,
-            spatial_size=tuple(input_size[1:]),
+            spatial_size=(input_size[1], input_size[2]),
             temporal_size=input_size[0],
             rel_pos=rel_pos,
         )