Fix type annotations

cvnets/models/classification/mobilenetv1.py

@@ -4,7 +4,7 @@
 #
 
 import argparse
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from torch import nn
 
@@ -143,7 +143,7 @@ def add_arguments(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentPars
     def _make_layer(
         self,
         opts,
-        mv1_config: Dict or List,
+        mv1_config: Union[Dict, List],
         input_channel: int,
         dilate: Optional[bool] = False,
         *args,

cvnets/models/classification/mobilenetv2.py

@@ -4,7 +4,7 @@
 #
 
 import argparse
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from torch import nn
 
@@ -159,7 +159,7 @@ def add_arguments(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentPars
     def _make_layer(
         self,
         opts,
-        mv2_config: Dict or List,
+        mv2_config: Union[Dict, List],
         width_mult: float,
         input_channel: int,
         dilate: Optional[bool] = False,

cvnets/models/segmentation/heads/base_seg_head.py

@@ -4,7 +4,7 @@
 #
 
 import argparse
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Union
 
 from torch import Tensor, nn
 
@@ -97,7 +97,7 @@ def forward_aux_head(self, enc_out: Dict) -> Tensor:
     def forward_seg_head(self, enc_out: Dict) -> Tensor:
         raise NotImplementedError
 
-    def forward(self, enc_out: Dict, *args, **kwargs) -> Tensor or Tuple[Tensor]:
+    def forward(self, enc_out: Dict, *args, **kwargs) -> Union[Tensor, Tuple[Tensor]]:
         out = self.forward_seg_head(enc_out=enc_out)
 
         if self.upsample_seg_out is not None:

data/transforms/image_pil.py

@@ -90,7 +90,7 @@ def _crop_fn(data: Dict, top: int, left: int, height: int, width: int) -> Dict:
 def _resize_fn(
     data: Dict,
     size: Union[Sequence, int],
-    interpolation: Optional[T.InterpolationMode or str] = T.InterpolationMode.BILINEAR,
+    interpolation: Optional[Union[T.InterpolationMode, str]] = T.InterpolationMode.BILINEAR,
 ) -> Dict:
     """Helper function for resizing"""
     img = data["image"]

data/transforms/video.py

@@ -716,7 +716,7 @@ class CenterCrop(BaseTransformation):
         size.
     """
 
-    def __init__(self, opts, size: Sequence or int, *args, **kwargs) -> None:
+    def __init__(self, opts, size: Union[Sequence, int], *args, **kwargs) -> None:
         super().__init__(opts=opts)
         if isinstance(size, Sequence) and len(size) == 2:
             self.height, self.width = size[0], size[1]

metrics/stats.py

@@ -166,7 +166,7 @@ def iter_summary(
         n_processed_samples: int,
         total_samples: int,
         elapsed_time: float,
-        learning_rate: float or list,
+        learning_rate: Union[float, list],
     ) -> None:
         if self.is_master_node:
             metric_stats = self._avg_statistics_all()