From b2f6615a59c1047a91affef2ed893d38839c7e27 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Thu, 20 May 2021 12:47:27 +0530
Subject: [PATCH 01/13] fix

---
 docs/Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/Makefile b/docs/Makefile
index 58daa471f5c..ef4cdc0cee0 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -2,7 +2,7 @@
 #
 
 # You can set these variables from the command line.
-SPHINXOPTS    = -W  # turn warnings into errors
+SPHINXOPTS    =   # turn warnings into errors
 SPHINXBUILD   = sphinx-build
 SPHINXPROJ    = torchvision
 SOURCEDIR     = source

From 304fa9e546da698efa6c8101f6ce2a9ab3f45b78 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Fri, 30 Jul 2021 17:11:40 +0530
Subject: [PATCH 02/13] add typings

---
 torchvision/models/quantization/googlenet.py  | 58 +++++++++---
 torchvision/models/quantization/inception.py  | 92 +++++++++++++++----
 .../models/quantization/mobilenetv2.py        | 33 +++++--
 .../models/quantization/mobilenetv3.py        | 37 ++++++--
 torchvision/models/quantization/resnet.py     | 64 ++++++++++---
 .../models/quantization/shufflenetv2.py       | 59 ++++++++++--
 torchvision/models/quantization/utils.py      |  4 +-
 7 files changed, 280 insertions(+), 67 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index bc1477d8f65..030922d15df 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -2,6 +2,8 @@
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
+from typing import Any
+from torch import Tensor
 
 from ..._internally_replaced_utils import load_state_dict_from_url
 from torchvision.models.googlenet import (
@@ -18,7 +20,13 @@
 }
 
 
-def googlenet(pretrained=False, progress=True, quantize=False, **kwargs):
+def googlenet(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
+
     r"""GoogLeNet (Inception v1) model architecture from
     `"Going Deeper with Convolutions" <http://arxiv.org/abs/1409.4842>`_.
 
@@ -77,41 +85,62 @@ def googlenet(pretrained=False, progress=True, quantize=False, **kwargs):
 
 class QuantizableBasicConv2d(BasicConv2d):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableBasicConv2d, self).__init__(*args, **kwargs)
         self.relu = nn.ReLU()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor
+    ) -> Tensor:
         x = self.conv(x)
         x = self.bn(x)
         x = self.relu(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
 
 class QuantizableInception(Inception):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInception, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ) -> Tensor:
         outputs = self._forward(x)
         return self.cat.cat(outputs, 1)
 
 
 class QuantizableInceptionAux(InceptionAux):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionAux, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(0.7)
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ) -> Tensor:
         # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14
         x = F.adaptive_avg_pool2d(x, (4, 4))
         # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4
@@ -131,7 +160,11 @@ def forward(self, x):
 
 class QuantizableGoogLeNet(GoogLeNet):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
         super(QuantizableGoogLeNet, self).__init__(
             blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux],
             *args,
@@ -140,7 +173,10 @@ def __init__(self, *args, **kwargs):
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x,
+    ):
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux1, aux2 = self._forward(x)
@@ -153,7 +189,7 @@ def forward(self, x):
         else:
             return self.eager_outputs(x, aux2, aux1)
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         r"""Fuse conv/bn/relu modules in googlenet model
 
         Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization.
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index 833d8fb8b75..a28840fe0d7 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -3,6 +3,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch import Tensor
+from typing import Any
+
 from torchvision.models import inception as inception_module
 from torchvision.models.inception import InceptionOutputs
 from ..._internally_replaced_utils import load_state_dict_from_url
@@ -22,7 +25,12 @@
 }
 
 
-def inception_v3(pretrained=False, progress=True, quantize=False, **kwargs):
+def inception_v3(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any,
+):
     r"""Inception v3 model architecture from
     `"Rethinking the Inception Architecture for Computer Vision" <http://arxiv.org/abs/1512.00567>`_.
 
@@ -84,32 +92,50 @@ def inception_v3(pretrained=False, progress=True, quantize=False, **kwargs):
 
 
 class QuantizableBasicConv2d(inception_module.BasicConv2d):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableBasicConv2d, self).__init__(*args, **kwargs)
         self.relu = nn.ReLU()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ) -> Tensor:
         x = self.conv(x)
         x = self.bn(x)
         x = self.relu(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
 
 class QuantizableInceptionA(inception_module.InceptionA):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionA, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ):
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
 
 class QuantizableInceptionB(inception_module.InceptionB):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionB, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
@@ -119,33 +145,48 @@ def forward(self, x):
 
 
 class QuantizableInceptionC(inception_module.InceptionC):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
         super(QuantizableInceptionC, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
 
 class QuantizableInceptionD(inception_module.InceptionD):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionD, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
 
 class QuantizableInceptionE(inception_module.InceptionE):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionE, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop1 = nn.quantized.FloatFunctional()
         self.myop2 = nn.quantized.FloatFunctional()
         self.myop3 = nn.quantized.FloatFunctional()
 
-    def _forward(self, x):
+    def _forward(
+        self,
+        x,
+    ):
         branch1x1 = self.branch1x1(x)
 
         branch3x3 = self.branch3x3_1(x)
@@ -166,18 +207,30 @@ def _forward(self, x):
         outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
         return outputs
 
-    def forward(self, x):
+    def forward(
+        self,
+        x,
+    ):
         outputs = self._forward(x)
         return self.myop3.cat(outputs, 1)
 
 
 class QuantizableInceptionAux(inception_module.InceptionAux):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInceptionAux, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
 
 
 class QuantizableInception3(inception_module.Inception3):
-    def __init__(self, num_classes=1000, aux_logits=True, transform_input=False):
+    def __init__(
+        self,
+        num_classes: int = 1000,
+        aux_logits: bool = True,
+        transform_input: bool = False,
+    ) -> None:
         super(QuantizableInception3, self).__init__(
             num_classes=num_classes,
             aux_logits=aux_logits,
@@ -195,7 +248,10 @@ def __init__(self, num_classes=1000, aux_logits=True, transform_input=False):
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x,
+    ):
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux = self._forward(x)
@@ -208,7 +264,7 @@ def forward(self, x):
         else:
             return self.eager_outputs(x, aux)
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         r"""Fuse conv/bn/relu modules in inception model
 
         Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization.
diff --git a/torchvision/models/quantization/mobilenetv2.py b/torchvision/models/quantization/mobilenetv2.py
index 857d919b1fa..2f4919b33c9 100644
--- a/torchvision/models/quantization/mobilenetv2.py
+++ b/torchvision/models/quantization/mobilenetv2.py
@@ -1,5 +1,10 @@
 from torch import nn
+from torch import Tensor
+
 from ..._internally_replaced_utils import load_state_dict_from_url
+
+from typing import Any
+
 from torchvision.models.mobilenetv2 import InvertedResidual, ConvBNReLU, MobileNetV2, model_urls
 from torch.quantization import QuantStub, DeQuantStub, fuse_modules
 from .utils import _replace_relu, quantize_model
@@ -14,7 +19,11 @@
 
 
 class QuantizableInvertedResidual(InvertedResidual):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
@@ -24,14 +33,18 @@ def forward(self, x):
         else:
             return self.conv(x)
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         for idx in range(len(self.conv)):
             if type(self.conv[idx]) == nn.Conv2d:
                 fuse_modules(self.conv, [str(idx), str(idx + 1)], inplace=True)
 
 
 class QuantizableMobileNetV2(MobileNetV2):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         """
         MobileNet V2 main class
 
@@ -42,13 +55,16 @@ def __init__(self, *args, **kwargs):
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ):
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         for m in self.modules():
             if type(m) == ConvBNReLU:
                 fuse_modules(m, ['0', '1', '2'], inplace=True)
@@ -56,7 +72,12 @@ def fuse_model(self):
                 m.fuse_model()
 
 
-def mobilenet_v2(pretrained=False, progress=True, quantize=False, **kwargs):
+def mobilenet_v2(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any,
+) -> QuantizableMobileNetV2:
     """
     Constructs a MobileNetV2 architecture from
     `"MobileNetV2: Inverted Residuals and Linear Bottlenecks"
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index 5462af89127..bc0b3fe8a8d 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -17,19 +17,27 @@
 
 
 class QuantizableSqueezeExcitation(SqueezeExcitation):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super().__init__(*args, **kwargs)
         self.skip_mul = nn.quantized.FloatFunctional()
 
     def forward(self, input: Tensor) -> Tensor:
         return self.skip_mul.mul(self._scale(input, False), input)
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         fuse_modules(self, ['fc1', 'relu'], inplace=True)
 
 
 class QuantizableInvertedResidual(InvertedResidual):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super().__init__(*args, se_layer=QuantizableSqueezeExcitation, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
@@ -41,7 +49,11 @@ def forward(self, x):
 
 
 class QuantizableMobileNetV3(MobileNetV3):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         """
         MobileNet V3 main class
 
@@ -52,13 +64,16 @@ def __init__(self, *args, **kwargs):
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: Tensor,
+    ):
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         for m in self.modules():
             if type(m) == ConvBNActivation:
                 modules_to_fuse = ['0', '1']
@@ -89,7 +104,8 @@ def _mobilenet_v3_model(
     progress: bool,
     quantize: bool,
     **kwargs: Any
-):
+) -> QuantizableMobileNetV3:
+
     model = QuantizableMobileNetV3(inverted_residual_setting, last_channel, block=QuantizableInvertedResidual, **kwargs)
     _replace_relu(model)
 
@@ -112,7 +128,12 @@ def _mobilenet_v3_model(
     return model
 
 
-def mobilenet_v3_large(pretrained=False, progress=True, quantize=False, **kwargs):
+def mobilenet_v3_large(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any,
+) -> QuantizableMobileNetV3:
     """
     Constructs a MobileNetV3 Large architecture from
     `"Searching for MobileNetV3" <https://arxiv.org/abs/1905.02244>`_.
diff --git a/torchvision/models/quantization/resnet.py b/torchvision/models/quantization/resnet.py
index 2f3f50e8013..71d94930b04 100644
--- a/torchvision/models/quantization/resnet.py
+++ b/torchvision/models/quantization/resnet.py
@@ -1,6 +1,9 @@
 import torch
 from torchvision.models.resnet import Bottleneck, BasicBlock, ResNet, model_urls
 import torch.nn as nn
+from torch import Tensor
+from typing import Any, Type, Union, List
+
 from ..._internally_replaced_utils import load_state_dict_from_url
 from torch.quantization import fuse_modules
 from .utils import _replace_relu, quantize_model
@@ -20,11 +23,15 @@
 
 
 class QuantizableBasicBlock(BasicBlock):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableBasicBlock, self).__init__(*args, **kwargs)
         self.add_relu = torch.nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         identity = x
 
         out = self.conv1(x)
@@ -41,7 +48,7 @@ def forward(self, x):
 
         return out
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         torch.quantization.fuse_modules(self, [['conv1', 'bn1', 'relu'],
                                                ['conv2', 'bn2']], inplace=True)
         if self.downsample:
@@ -49,13 +56,17 @@ def fuse_model(self):
 
 
 class QuantizableBottleneck(Bottleneck):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
         super(QuantizableBottleneck, self).__init__(*args, **kwargs)
         self.skip_add_relu = nn.quantized.FloatFunctional()
         self.relu1 = nn.ReLU(inplace=False)
         self.relu2 = nn.ReLU(inplace=False)
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         identity = x
         out = self.conv1(x)
         out = self.bn1(out)
@@ -83,13 +94,17 @@ def fuse_model(self):
 
 class QuantizableResNet(ResNet):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
         super(QuantizableResNet, self).__init__(*args, **kwargs)
 
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         x = self.quant(x)
         # Ensure scriptability
         # super(QuantizableResNet,self).forward(x)
@@ -98,7 +113,7 @@ def forward(self, x):
         x = self.dequant(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         r"""Fuse conv/bn/relu modules in resnet models
 
         Fuse conv+bn+relu/ Conv+relu/conv+Bn modules to prepare for quantization.
@@ -112,7 +127,16 @@ def fuse_model(self):
                 m.fuse_model()
 
 
-def _resnet(arch, block, layers, pretrained, progress, quantize, **kwargs):
+def _resnet(
+    arch: str,
+    block: Type[Union[BasicBlock, Bottleneck]],
+    layers: List[int],
+    pretrained: bool,
+    progress: bool,
+    quantize: bool,
+    **kwargs: Any,
+) -> QuantizableResNet:
+
     model = QuantizableResNet(block, layers, **kwargs)
     _replace_relu(model)
     if quantize:
@@ -135,7 +159,12 @@ def _resnet(arch, block, layers, pretrained, progress, quantize, **kwargs):
     return model
 
 
-def resnet18(pretrained=False, progress=True, quantize=False, **kwargs):
+def resnet18(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     r"""ResNet-18 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
@@ -148,7 +177,13 @@ def resnet18(pretrained=False, progress=True, quantize=False, **kwargs):
                    quantize, **kwargs)
 
 
-def resnet50(pretrained=False, progress=True, quantize=False, **kwargs):
+def resnet50(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
+
     r"""ResNet-50 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
@@ -161,7 +196,12 @@ def resnet50(pretrained=False, progress=True, quantize=False, **kwargs):
                    quantize, **kwargs)
 
 
-def resnext101_32x8d(pretrained=False, progress=True, quantize=False, **kwargs):
+def resnext101_32x8d(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     r"""ResNeXt-101 32x8d model from
     `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
 
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 17885015772..5f00b3d1635 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -1,5 +1,8 @@
 import torch
 import torch.nn as nn
+from torch import Tensor
+from typing import Any
+
 from ..._internally_replaced_utils import load_state_dict_from_url
 import torchvision.models.shufflenetv2
 import sys
@@ -22,11 +25,15 @@
 
 
 class QuantizableInvertedResidual(shufflenetv2.InvertedResidual):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
             out = self.cat.cat((x1, self.branch2(x2)), dim=1)
@@ -39,18 +46,22 @@ def forward(self, x):
 
 
 class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         super(QuantizableShuffleNetV2, self).__init__(*args, inverted_residual=QuantizableInvertedResidual, **kwargs)
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
         return x
 
-    def fuse_model(self):
+    def fuse_model(self) -> Tensor:
         r"""Fuse conv/bn/relu modules in shufflenetv2 model
 
         Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization.
@@ -74,7 +85,15 @@ def fuse_model(self):
                 )
 
 
-def _shufflenetv2(arch, pretrained, progress, quantize, *args, **kwargs):
+def _shufflenetv2(
+    arch: str,
+    pretrained: bool,
+    progress: bool,
+    quantize: bool,
+    *args: Any,
+    **kwargs: Any,
+) -> QuantizableShuffleNetV2:
+
     model = QuantizableShuffleNetV2(*args, **kwargs)
     _replace_relu(model)
 
@@ -98,7 +117,12 @@ def _shufflenetv2(arch, pretrained, progress, quantize, *args, **kwargs):
     return model
 
 
-def shufflenet_v2_x0_5(pretrained=False, progress=True, quantize=False, **kwargs):
+def shufflenet_v2_x0_5(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     """
     Constructs a ShuffleNetV2 with 0.5x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -113,7 +137,12 @@ def shufflenet_v2_x0_5(pretrained=False, progress=True, quantize=False, **kwargs
                          [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs)
 
 
-def shufflenet_v2_x1_0(pretrained=False, progress=True, quantize=False, **kwargs):
+def shufflenet_v2_x1_0(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     """
     Constructs a ShuffleNetV2 with 1.0x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -128,7 +157,12 @@ def shufflenet_v2_x1_0(pretrained=False, progress=True, quantize=False, **kwargs
                          [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)
 
 
-def shufflenet_v2_x1_5(pretrained=False, progress=True, quantize=False, **kwargs):
+def shufflenet_v2_x1_5(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     """
     Constructs a ShuffleNetV2 with 1.5x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -143,7 +177,12 @@ def shufflenet_v2_x1_5(pretrained=False, progress=True, quantize=False, **kwargs
                          [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)
 
 
-def shufflenet_v2_x2_0(pretrained=False, progress=True, quantize=False, **kwargs):
+def shufflenet_v2_x2_0(
+    pretrained: bool = False,
+    progress: bool = True,
+    quantize: bool = False,
+    **kwargs: Any
+):
     """
     Constructs a ShuffleNetV2 with 2.0x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py
index bf23c9a9332..a7a81c1e167 100644
--- a/torchvision/models/quantization/utils.py
+++ b/torchvision/models/quantization/utils.py
@@ -2,7 +2,7 @@
 from torch import nn
 
 
-def _replace_relu(module):
+def _replace_relu(module: nn.Module) -> None:
     reassign = {}
     for name, mod in module.named_children():
         _replace_relu(mod)
@@ -16,7 +16,7 @@ def _replace_relu(module):
         module._modules[key] = value
 
 
-def quantize_model(model, backend):
+def quantize_model(model: nn.Module, backend: str) -> None:
     _dummy_input_data = torch.rand(1, 3, 299, 299)
     if backend not in torch.backends.quantized.supported_engines:
         raise RuntimeError("Quantized backend not supported ")

From a20ee80ea2e5f8ffb5203ebd65c903e7e0aa4849 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Fri, 30 Jul 2021 19:14:57 +0530
Subject: [PATCH 03/13] fixup some more types

---
 torchvision/models/quantization/googlenet.py  | 44 +++---------
 torchvision/models/quantization/inception.py  | 69 ++++---------------
 .../models/quantization/mobilenetv2.py        | 19 ++---
 .../models/quantization/mobilenetv3.py        | 25 ++-----
 torchvision/models/quantization/resnet.py     | 20 ++----
 .../models/quantization/shufflenetv2.py       | 12 +---
 6 files changed, 36 insertions(+), 153 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index 030922d15df..96b33b8ac2d 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -85,18 +85,11 @@ def googlenet(
 
 class QuantizableBasicConv2d(BasicConv2d):
 
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any,) -> None:
         super(QuantizableBasicConv2d, self).__init__(*args, **kwargs)
         self.relu = nn.ReLU()
 
-    def forward(
-        self,
-        x: Tensor
-    ) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv(x)
         x = self.bn(x)
         x = self.relu(x)
@@ -108,39 +101,25 @@ def fuse_model(self) -> None:
 
 class QuantizableInception(Inception):
 
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any,) -> None:
         super(QuantizableInception, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(
-        self,
-        x: Tensor,
-    ) -> Tensor:
+    def forward(self, x: Tensor,) -> Tensor:
         outputs = self._forward(x)
         return self.cat.cat(outputs, 1)
 
 
 class QuantizableInceptionAux(InceptionAux):
 
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any,) -> None:
         super(QuantizableInceptionAux, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(0.7)
 
-    def forward(
-        self,
-        x: Tensor,
-    ) -> Tensor:
+    def forward(self, x: Tensor,) -> Tensor:
         # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14
         x = F.adaptive_avg_pool2d(x, (4, 4))
         # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4
@@ -160,11 +139,7 @@ def forward(
 
 class QuantizableGoogLeNet(GoogLeNet):
 
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableGoogLeNet, self).__init__(
             blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux],
             *args,
@@ -173,10 +148,7 @@ def __init__(
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(
-        self,
-        x,
-    ):
+    def forward(self, x: Tensor):
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux1, aux2 = self._forward(x)
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index a28840fe0d7..319a880cf83 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -92,18 +92,11 @@ def inception_v3(
 
 
 class QuantizableBasicConv2d(inception_module.BasicConv2d):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableBasicConv2d, self).__init__(*args, **kwargs)
         self.relu = nn.ReLU()
 
-    def forward(
-        self,
-        x: Tensor,
-    ) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv(x)
         x = self.bn(x)
         x = self.relu(x)
@@ -114,42 +107,27 @@ def fuse_model(self) -> None:
 
 
 class QuantizableInceptionA(inception_module.InceptionA):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionA, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(
-        self,
-        x: Tensor,
-    ):
+    def forward(self, x: Tensor):
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
 
 class QuantizableInceptionB(inception_module.InceptionB):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionB, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
 
 class QuantizableInceptionC(inception_module.InceptionC):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionC, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
@@ -159,11 +137,7 @@ def forward(self, x: Tensor):
 
 
 class QuantizableInceptionD(inception_module.InceptionD):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionD, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
@@ -173,20 +147,13 @@ def forward(self, x: Tensor):
 
 
 class QuantizableInceptionE(inception_module.InceptionE):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionE, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop1 = nn.quantized.FloatFunctional()
         self.myop2 = nn.quantized.FloatFunctional()
         self.myop3 = nn.quantized.FloatFunctional()
 
-    def _forward(
-        self,
-        x,
-    ):
+    def _forward(self, x: Tensor):
         branch1x1 = self.branch1x1(x)
 
         branch3x3 = self.branch3x3_1(x)
@@ -207,20 +174,13 @@ def _forward(
         outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
         return outputs
 
-    def forward(
-        self,
-        x,
-    ):
+    def forward(self, x: Tensor):
         outputs = self._forward(x)
         return self.myop3.cat(outputs, 1)
 
 
 class QuantizableInceptionAux(inception_module.InceptionAux):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionAux, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
 
 
@@ -248,10 +208,7 @@ def __init__(
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(
-        self,
-        x,
-    ):
+    def forward(self, x: Tensor):
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux = self._forward(x)
diff --git a/torchvision/models/quantization/mobilenetv2.py b/torchvision/models/quantization/mobilenetv2.py
index 2f4919b33c9..ceec1719e5f 100644
--- a/torchvision/models/quantization/mobilenetv2.py
+++ b/torchvision/models/quantization/mobilenetv2.py
@@ -19,15 +19,11 @@
 
 
 class QuantizableInvertedResidual(InvertedResidual):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         if self.use_res_connect:
             return self.skip_add.add(x, self.conv(x))
         else:
@@ -40,11 +36,7 @@ def fuse_model(self) -> None:
 
 
 class QuantizableMobileNetV2(MobileNetV2):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """
         MobileNet V2 main class
 
@@ -55,10 +47,7 @@ def __init__(
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(
-        self,
-        x: Tensor,
-    ):
+    def forward(self, x: Tensor):
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index bc0b3fe8a8d..045aeb86099 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -17,11 +17,7 @@
 
 
 class QuantizableSqueezeExcitation(SqueezeExcitation):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.skip_mul = nn.quantized.FloatFunctional()
 
@@ -33,15 +29,11 @@ def fuse_model(self) -> None:
 
 
 class QuantizableInvertedResidual(InvertedResidual):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, se_layer=QuantizableSqueezeExcitation, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         if self.use_res_connect:
             return self.skip_add.add(x, self.block(x))
         else:
@@ -49,11 +41,7 @@ def forward(self, x):
 
 
 class QuantizableMobileNetV3(MobileNetV3):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """
         MobileNet V3 main class
 
@@ -64,10 +52,7 @@ def __init__(
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(
-        self,
-        x: Tensor,
-    ):
+    def forward(self, x: Tensor):
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
diff --git a/torchvision/models/quantization/resnet.py b/torchvision/models/quantization/resnet.py
index 71d94930b04..c9b92c0f175 100644
--- a/torchvision/models/quantization/resnet.py
+++ b/torchvision/models/quantization/resnet.py
@@ -23,11 +23,7 @@
 
 
 class QuantizableBasicBlock(BasicBlock):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableBasicBlock, self).__init__(*args, **kwargs)
         self.add_relu = torch.nn.quantized.FloatFunctional()
 
@@ -56,11 +52,7 @@ def fuse_model(self) -> None:
 
 
 class QuantizableBottleneck(Bottleneck):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableBottleneck, self).__init__(*args, **kwargs)
         self.skip_add_relu = nn.quantized.FloatFunctional()
         self.relu1 = nn.ReLU(inplace=False)
@@ -84,7 +76,7 @@ def forward(self, x: Tensor):
 
         return out
 
-    def fuse_model(self):
+    def fuse_model(self) -> None:
         fuse_modules(self, [['conv1', 'bn1', 'relu1'],
                             ['conv2', 'bn2', 'relu2'],
                             ['conv3', 'bn3']], inplace=True)
@@ -94,11 +86,7 @@ def fuse_model(self):
 
 class QuantizableResNet(ResNet):
 
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableResNet, self).__init__(*args, **kwargs)
 
         self.quant = torch.quantization.QuantStub()
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 5f00b3d1635..5cc1ca95880 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -25,11 +25,7 @@
 
 
 class QuantizableInvertedResidual(shufflenetv2.InvertedResidual):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
@@ -46,11 +42,7 @@ def forward(self, x: Tensor):
 
 
 class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
-    def __init__(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableShuffleNetV2, self).__init__(*args, inverted_residual=QuantizableInvertedResidual, **kwargs)
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()

From 78ddc541ec1daaeed9edf92b3048ebf47203651d Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Fri, 30 Jul 2021 20:27:31 +0530
Subject: [PATCH 04/13] Type more

---
 torchvision/models/quantization/googlenet.py    | 14 +++++++-------
 torchvision/models/quantization/inception.py    | 16 ++++++++--------
 torchvision/models/quantization/mobilenetv2.py  |  4 ++--
 torchvision/models/quantization/mobilenetv3.py  |  6 +++---
 torchvision/models/quantization/resnet.py       | 12 ++++++------
 torchvision/models/quantization/shufflenetv2.py | 12 ++++++------
 6 files changed, 32 insertions(+), 32 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index 96b33b8ac2d..9d625fd529f 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -24,7 +24,7 @@ def googlenet(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
 
     r"""GoogLeNet (Inception v1) model architecture from
@@ -85,7 +85,7 @@ def googlenet(
 
 class QuantizableBasicConv2d(BasicConv2d):
 
-    def __init__(self, *args: Any, **kwargs: Any,) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableBasicConv2d, self).__init__(*args, **kwargs)
         self.relu = nn.ReLU()
 
@@ -101,25 +101,25 @@ def fuse_model(self) -> None:
 
 class QuantizableInception(Inception):
 
-    def __init__(self, *args: Any, **kwargs: Any,) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInception, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor,) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.cat.cat(outputs, 1)
 
 
 class QuantizableInceptionAux(InceptionAux):
 
-    def __init__(self, *args: Any, **kwargs: Any,) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionAux, self).__init__(
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(0.7)
 
-    def forward(self, x: Tensor,) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14
         x = F.adaptive_avg_pool2d(x, (4, 4))
         # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4
@@ -148,7 +148,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> GoogLeNetOutputs:
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux1, aux2 = self._forward(x)
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index 319a880cf83..fd173446cf0 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -4,7 +4,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
-from typing import Any
+from typing import Any, List
 
 from torchvision.models import inception as inception_module
 from torchvision.models.inception import InceptionOutputs
@@ -111,7 +111,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionA, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
@@ -121,7 +121,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionB, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
@@ -131,7 +131,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionC, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
@@ -141,7 +141,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionD, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.myop = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.myop.cat(outputs, 1)
 
@@ -153,7 +153,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.myop2 = nn.quantized.FloatFunctional()
         self.myop3 = nn.quantized.FloatFunctional()
 
-    def _forward(self, x: Tensor):
+    def _forward(self, x: Tensor) -> List[Tensor]:
         branch1x1 = self.branch1x1(x)
 
         branch3x3 = self.branch3x3_1(x)
@@ -174,7 +174,7 @@ def _forward(self, x: Tensor):
         outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
         return outputs
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         outputs = self._forward(x)
         return self.myop3.cat(outputs, 1)
 
@@ -208,7 +208,7 @@ def __init__(
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> InceptionOutputs:
         x = self._transform_input(x)
         x = self.quant(x)
         x, aux = self._forward(x)
diff --git a/torchvision/models/quantization/mobilenetv2.py b/torchvision/models/quantization/mobilenetv2.py
index ceec1719e5f..f914fdc7815 100644
--- a/torchvision/models/quantization/mobilenetv2.py
+++ b/torchvision/models/quantization/mobilenetv2.py
@@ -23,7 +23,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         if self.use_res_connect:
             return self.skip_add.add(x, self.conv(x))
         else:
@@ -47,7 +47,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index 045aeb86099..1766dc1c033 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -33,7 +33,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, se_layer=QuantizableSqueezeExcitation, **kwargs)
         self.skip_add = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         if self.use_res_connect:
             return self.skip_add.add(x, self.block(x))
         else:
@@ -52,7 +52,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.quant = QuantStub()
         self.dequant = DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
@@ -88,7 +88,7 @@ def _mobilenet_v3_model(
     pretrained: bool,
     progress: bool,
     quantize: bool,
-    **kwargs: Any
+    **kwargs: Any,
 ) -> QuantizableMobileNetV3:
 
     model = QuantizableMobileNetV3(inverted_residual_setting, last_channel, block=QuantizableInvertedResidual, **kwargs)
diff --git a/torchvision/models/quantization/resnet.py b/torchvision/models/quantization/resnet.py
index c9b92c0f175..634f95a6b4d 100644
--- a/torchvision/models/quantization/resnet.py
+++ b/torchvision/models/quantization/resnet.py
@@ -27,7 +27,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableBasicBlock, self).__init__(*args, **kwargs)
         self.add_relu = torch.nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         identity = x
 
         out = self.conv1(x)
@@ -58,7 +58,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.relu1 = nn.ReLU(inplace=False)
         self.relu2 = nn.ReLU(inplace=False)
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         identity = x
         out = self.conv1(x)
         out = self.bn1(out)
@@ -92,7 +92,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)
         # Ensure scriptability
         # super(QuantizableResNet,self).forward(x)
@@ -151,7 +151,7 @@ def resnet18(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     r"""ResNet-18 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
@@ -169,7 +169,7 @@ def resnet50(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
 
     r"""ResNet-50 model from
@@ -188,7 +188,7 @@ def resnext101_32x8d(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     r"""ResNeXt-101 32x8d model from
     `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 5cc1ca95880..6663388de9a 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -29,7 +29,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
             out = self.cat.cat((x1, self.branch2(x2)), dim=1)
@@ -47,7 +47,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 
-    def forward(self, x: Tensor):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.quant(x)
         x = self._forward_impl(x)
         x = self.dequant(x)
@@ -113,7 +113,7 @@ def shufflenet_v2_x0_5(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     """
     Constructs a ShuffleNetV2 with 0.5x output channels, as described in
@@ -133,7 +133,7 @@ def shufflenet_v2_x1_0(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     """
     Constructs a ShuffleNetV2 with 1.0x output channels, as described in
@@ -153,7 +153,7 @@ def shufflenet_v2_x1_5(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     """
     Constructs a ShuffleNetV2 with 1.5x output channels, as described in
@@ -173,7 +173,7 @@ def shufflenet_v2_x2_0(
     pretrained: bool = False,
     progress: bool = True,
     quantize: bool = False,
-    **kwargs: Any
+    **kwargs: Any,
 ):
     """
     Constructs a ShuffleNetV2 with 2.0x output channels, as described in

From 1e0de0c30beba5014eaf03c79dcfd262a9d1c057 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Mon, 2 Aug 2021 11:50:36 +0530
Subject: [PATCH 05/13] remove mypy ignore

---
 mypy.ini | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index 040b52dfda4..7035c4a77ff 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -20,10 +20,6 @@ ignore_errors=True
 
 ignore_errors = True
 
-[mypy-torchvision.models.quantization.*]
-
-ignore_errors = True
-
 [mypy-torchvision.ops.*]
 
 ignore_errors = True

From 4bc3a809bc24689c799ffcee78daac9d97f480f5 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Mon, 2 Aug 2021 11:55:18 +0530
Subject: [PATCH 06/13] add missing typings

---
 torchvision/models/quantization/mobilenetv3.py  | 2 +-
 torchvision/models/quantization/resnet.py       | 6 +++---
 torchvision/models/quantization/shufflenetv2.py | 8 ++++----
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index 1766dc1c033..44b6512ca6f 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -74,7 +74,7 @@ def _load_weights(
     model: QuantizableMobileNetV3,
     model_url: Optional[str],
     progress: bool,
-):
+) -> None:
     if model_url is None:
         raise ValueError("No checkpoint is available for {}".format(arch))
     state_dict = load_state_dict_from_url(model_url, progress=progress)
diff --git a/torchvision/models/quantization/resnet.py b/torchvision/models/quantization/resnet.py
index 634f95a6b4d..8f87e40ec3d 100644
--- a/torchvision/models/quantization/resnet.py
+++ b/torchvision/models/quantization/resnet.py
@@ -152,7 +152,7 @@ def resnet18(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableResNet:
     r"""ResNet-18 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
@@ -170,7 +170,7 @@ def resnet50(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableResNet:
 
     r"""ResNet-50 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
@@ -189,7 +189,7 @@ def resnext101_32x8d(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableResNet:
     r"""ResNeXt-101 32x8d model from
     `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
 
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 6663388de9a..0a11c45b107 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -114,7 +114,7 @@ def shufflenet_v2_x0_5(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableShuffleNetV2:
     """
     Constructs a ShuffleNetV2 with 0.5x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -134,7 +134,7 @@ def shufflenet_v2_x1_0(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableShuffleNetV2:
     """
     Constructs a ShuffleNetV2 with 1.0x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -154,7 +154,7 @@ def shufflenet_v2_x1_5(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableShuffleNetV2:
     """
     Constructs a ShuffleNetV2 with 1.5x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
@@ -174,7 +174,7 @@ def shufflenet_v2_x2_0(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> QuantizableShuffleNetV2:
     """
     Constructs a ShuffleNetV2 with 2.0x output channels, as described in
     `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"

From 3c1bd67e8508571db1f9ce53a7a4385fa80f5d46 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Mon, 2 Aug 2021 16:10:28 +0530
Subject: [PATCH 07/13] fix a few mypy errors

---
 torchvision/models/quantization/shufflenetv2.py | 12 +++++++-----
 torchvision/models/quantization/utils.py        |  3 ++-
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 0a11c45b107..cd01073b225 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -1,14 +1,15 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
-from typing import Any
+from typing import Any, List
 
 from ..._internally_replaced_utils import load_state_dict_from_url
-import torchvision.models.shufflenetv2
-import sys
+# import torchvision.models.shufflenetv2
+from torchvision.models import shufflenetv2
+# import sys
 from .utils import _replace_relu, quantize_model
 
-shufflenetv2 = sys.modules['torchvision.models.shufflenetv2']
+# shufflenetv2 = sys.modules['torchvision.models.shufflenetv2']
 
 __all__ = [
     'QuantizableShuffleNetV2', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
@@ -29,7 +30,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: List[Tensor]) -> Tensor:
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
             out = self.cat.cat((x1, self.branch2(x2)), dim=1)
@@ -75,6 +76,7 @@ def fuse_model(self) -> Tensor:
                     [["0", "1", "2"], ["3", "4"], ["5", "6", "7"]],
                     inplace=True,
                 )
+        return
 
 
 def _shufflenetv2(
diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py
index a7a81c1e167..b3a95de9122 100644
--- a/torchvision/models/quantization/utils.py
+++ b/torchvision/models/quantization/utils.py
@@ -1,5 +1,6 @@
 import torch
 from torch import nn
+from torch import Tensor, Union
 
 
 def _replace_relu(module: nn.Module) -> None:
@@ -16,7 +17,7 @@ def _replace_relu(module: nn.Module) -> None:
         module._modules[key] = value
 
 
-def quantize_model(model: nn.Module, backend: str) -> None:
+def quantize_model(model: Union[nn.Module, Tensor], backend: str) -> None:
     _dummy_input_data = torch.rand(1, 3, 299, 299)
     if backend not in torch.backends.quantized.supported_engines:
         raise RuntimeError("Quantized backend not supported ")

From 479c64c9b814a9fca36efb1709a5019f085b07e8 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Mon, 2 Aug 2021 22:53:56 +0530
Subject: [PATCH 08/13] fix mypy errors

---
 torchvision/models/quantization/shufflenetv2.py | 11 +++++------
 torchvision/models/quantization/utils.py        |  9 ++++-----
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index cd01073b225..923949df509 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
-from typing import Any, List
+from typing import Any
 
 from ..._internally_replaced_utils import load_state_dict_from_url
 # import torchvision.models.shufflenetv2
@@ -30,12 +30,12 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInvertedResidual, self).__init__(*args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
-    def forward(self, x: List[Tensor]) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
-            out = self.cat.cat((x1, self.branch2(x2)), dim=1)
+            out = self.cat.cat((x1, self.branch2(x2)), dim=1)  # type: ignore
         else:
-            out = self.cat.cat((self.branch1(x), self.branch2(x)), dim=1)
+            out = self.cat.cat((self.branch1(x), self.branch2(x)), dim=1)  # type: ignore
 
         out = shufflenetv2.channel_shuffle(out, 2)
 
@@ -54,7 +54,7 @@ def forward(self, x: Tensor) -> Tensor:
         x = self.dequant(x)
         return x
 
-    def fuse_model(self) -> Tensor:
+    def fuse_model(self) -> None:
         r"""Fuse conv/bn/relu modules in shufflenetv2 model
 
         Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization.
@@ -76,7 +76,6 @@ def fuse_model(self) -> Tensor:
                     [["0", "1", "2"], ["3", "4"], ["5", "6", "7"]],
                     inplace=True,
                 )
-        return
 
 
 def _shufflenetv2(
diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py
index b3a95de9122..5b9b6ec2009 100644
--- a/torchvision/models/quantization/utils.py
+++ b/torchvision/models/quantization/utils.py
@@ -1,6 +1,5 @@
 import torch
 from torch import nn
-from torch import Tensor, Union
 
 
 def _replace_relu(module: nn.Module) -> None:
@@ -17,7 +16,7 @@ def _replace_relu(module: nn.Module) -> None:
         module._modules[key] = value
 
 
-def quantize_model(model: Union[nn.Module, Tensor], backend: str) -> None:
+def quantize_model(model: nn.Module, backend: str) -> None:
     _dummy_input_data = torch.rand(1, 3, 299, 299)
     if backend not in torch.backends.quantized.supported_engines:
         raise RuntimeError("Quantized backend not supported ")
@@ -25,15 +24,15 @@ def quantize_model(model: Union[nn.Module, Tensor], backend: str) -> None:
     model.eval()
     # Make sure that weight qconfig matches that of the serialized models
     if backend == 'fbgemm':
-        model.qconfig = torch.quantization.QConfig(
+        model.qconfig = torch.quantization.QConfig(  # type: ignore
             activation=torch.quantization.default_observer,
             weight=torch.quantization.default_per_channel_weight_observer)
     elif backend == 'qnnpack':
-        model.qconfig = torch.quantization.QConfig(
+        model.qconfig = torch.quantization.QConfig(  # type: ignore
             activation=torch.quantization.default_observer,
             weight=torch.quantization.default_weight_observer)
 
-    model.fuse_model()
+    model.fuse_model()  # type: ignore
     torch.quantization.prepare(model, inplace=True)
     model(_dummy_input_data)
     torch.quantization.convert(model, inplace=True)

From ee1b93e7cc046b7653d5bc7217d26f9d11b2eb32 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Mon, 2 Aug 2021 22:54:05 +0530
Subject: [PATCH 09/13] fix mypy

---
 torchvision/models/quantization/googlenet.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index 9d625fd529f..bc9ee78a048 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -78,8 +78,8 @@ def googlenet(
 
         if not original_aux_logits:
             model.aux_logits = False
-            model.aux1 = None
-            model.aux2 = None
+            model.aux1 = None  # type: ignore
+            model.aux2 = None  # type: ignore
     return model
 
 

From 73f28df4006fb021a98cdd5871ccbd115574221b Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Thu, 5 Aug 2021 10:59:05 +0530
Subject: [PATCH 10/13] ignore types

---
 torchvision/models/quantization/googlenet.py    | 4 ++--
 torchvision/models/quantization/shufflenetv2.py | 8 ++------
 torchvision/models/quantization/utils.py        | 6 +++---
 3 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index bc9ee78a048..d368d561f14 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -78,8 +78,8 @@ def googlenet(
 
         if not original_aux_logits:
             model.aux_logits = False
-            model.aux1 = None  # type: ignore
-            model.aux2 = None  # type: ignore
+            model.aux1 = None  # type: ignore[assignment]
+            model.aux2 = None  # type: ignore[assignment]
     return model
 
 
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 923949df509..4a0ce2d2881 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -4,13 +4,9 @@
 from typing import Any
 
 from ..._internally_replaced_utils import load_state_dict_from_url
-# import torchvision.models.shufflenetv2
 from torchvision.models import shufflenetv2
-# import sys
 from .utils import _replace_relu, quantize_model
 
-# shufflenetv2 = sys.modules['torchvision.models.shufflenetv2']
-
 __all__ = [
     'QuantizableShuffleNetV2', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
     'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'
@@ -33,9 +29,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     def forward(self, x: Tensor) -> Tensor:
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
-            out = self.cat.cat((x1, self.branch2(x2)), dim=1)  # type: ignore
+            out = self.cat.cat((x1, self.branch2(x2)), dim=1)  # type: ignore[arg-type]
         else:
-            out = self.cat.cat((self.branch1(x), self.branch2(x)), dim=1)  # type: ignore
+            out = self.cat.cat((self.branch1(x), self.branch2(x)), dim=1)  # type: ignore[arg-type]
 
         out = shufflenetv2.channel_shuffle(out, 2)
 
diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py
index 5b9b6ec2009..941db439e0c 100644
--- a/torchvision/models/quantization/utils.py
+++ b/torchvision/models/quantization/utils.py
@@ -24,15 +24,15 @@ def quantize_model(model: nn.Module, backend: str) -> None:
     model.eval()
     # Make sure that weight qconfig matches that of the serialized models
     if backend == 'fbgemm':
-        model.qconfig = torch.quantization.QConfig(  # type: ignore
+        model.qconfig = torch.quantization.QConfig(  # type: ignore[assignment]
             activation=torch.quantization.default_observer,
             weight=torch.quantization.default_per_channel_weight_observer)
     elif backend == 'qnnpack':
-        model.qconfig = torch.quantization.QConfig(  # type: ignore
+        model.qconfig = torch.quantization.QConfig(  # type: ignore[assignment]
             activation=torch.quantization.default_observer,
             weight=torch.quantization.default_weight_observer)
 
-    model.fuse_model()  # type: ignore
+    model.fuse_model()  # type: ignore[operator]
     torch.quantization.prepare(model, inplace=True)
     model(_dummy_input_data)
     torch.quantization.convert(model, inplace=True)

From 2528d13aab56050db7eda5dc1119c17fb2a1de21 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Tue, 17 Aug 2021 01:24:24 +0530
Subject: [PATCH 11/13] fixup annotation

---
 torchvision/models/quantization/googlenet.py | 2 +-
 torchvision/models/quantization/inception.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index d368d561f14..3b1779c511e 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -25,7 +25,7 @@ def googlenet(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> "QuantizableGoogLeNet":
 
     r"""GoogLeNet (Inception v1) model architecture from
     `"Going Deeper with Convolutions" <http://arxiv.org/abs/1409.4842>`_.
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index fd173446cf0..91862dda4e1 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -30,7 +30,8 @@ def inception_v3(
     progress: bool = True,
     quantize: bool = False,
     **kwargs: Any,
-):
+) -> "QuantizableInception3":
+
     r"""Inception v3 model architecture from
     `"Rethinking the Inception Architecture for Computer Vision" <http://arxiv.org/abs/1512.00567>`_.
 

From 2bba20771c58106eff03a607e122f5f3569f12f1 Mon Sep 17 00:00:00 2001
From: Aditya Oke <okeaditya315@gmail.com>
Date: Wed, 18 Aug 2021 00:28:34 +0530
Subject: [PATCH 12/13] fix remaining types

---
 torchvision/models/quantization/googlenet.py  | 15 +++++---
 torchvision/models/quantization/inception.py  | 38 ++++++++++++++++---
 .../models/quantization/mobilenetv3.py        |  7 +++-
 .../models/quantization/shufflenetv2.py       | 11 ++++--
 4 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py
index 3b1779c511e..685815ac676 100644
--- a/torchvision/models/quantization/googlenet.py
+++ b/torchvision/models/quantization/googlenet.py
@@ -102,7 +102,7 @@ def fuse_model(self) -> None:
 class QuantizableInception(Inception):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInception, self).__init__(
+        super(QuantizableInception, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d, *args, **kwargs)
         self.cat = nn.quantized.FloatFunctional()
 
@@ -112,10 +112,13 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionAux(InceptionAux):
-
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionAux, self).__init__(
-            conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionAux, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(0.7)
 
@@ -138,9 +141,9 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableGoogLeNet(GoogLeNet):
-
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableGoogLeNet, self).__init__(
+        super(QuantizableGoogLeNet, self).__init__(  # type: ignore[misc]
             blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux],
             *args,
             **kwargs
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index 91862dda4e1..71bcfc7b4fd 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -106,10 +106,16 @@ def forward(self, x: Tensor) -> Tensor:
     def fuse_model(self) -> None:
         torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
+# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
+
 
 class QuantizableInceptionA(inception_module.InceptionA):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionA, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionA, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.myop = nn.quantized.FloatFunctional()
 
     def forward(self, x: Tensor) -> Tensor:
@@ -119,7 +125,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableInceptionB(inception_module.InceptionB):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionB, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionB, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.myop = nn.quantized.FloatFunctional()
 
     def forward(self, x: Tensor) -> Tensor:
@@ -129,7 +139,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableInceptionC(inception_module.InceptionC):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionC, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionC, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.myop = nn.quantized.FloatFunctional()
 
     def forward(self, x: Tensor) -> Tensor:
@@ -139,7 +153,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableInceptionD(inception_module.InceptionD):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionD, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionD, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.myop = nn.quantized.FloatFunctional()
 
     def forward(self, x: Tensor) -> Tensor:
@@ -149,7 +167,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableInceptionE(inception_module.InceptionE):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionE, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionE, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
         self.myop1 = nn.quantized.FloatFunctional()
         self.myop2 = nn.quantized.FloatFunctional()
         self.myop3 = nn.quantized.FloatFunctional()
@@ -182,7 +204,11 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableInceptionAux(inception_module.InceptionAux):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableInceptionAux, self).__init__(conv_block=QuantizableBasicConv2d, *args, **kwargs)
+        super(QuantizableInceptionAux, self).__init__(  # type: ignore[misc]
+            conv_block=QuantizableBasicConv2d,
+            *args,
+            **kwargs
+        )
 
 
 class QuantizableInception3(inception_module.Inception3):
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index 44b6512ca6f..050db77001e 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -28,9 +28,14 @@ def fuse_model(self) -> None:
         fuse_modules(self, ['fc1', 'relu'], inplace=True)
 
 
+# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
 class QuantizableInvertedResidual(InvertedResidual):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, se_layer=QuantizableSqueezeExcitation, **kwargs)
+        super().__init__(  # type: ignore[misc]
+            se_layer=QuantizableSqueezeExcitation,
+            *args,
+            **kwargs
+        )
         self.skip_add = nn.quantized.FloatFunctional()
 
     def forward(self, x: Tensor) -> Tensor:
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 4a0ce2d2881..16be4f21ce5 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -29,9 +29,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     def forward(self, x: Tensor) -> Tensor:
         if self.stride == 1:
             x1, x2 = x.chunk(2, dim=1)
-            out = self.cat.cat((x1, self.branch2(x2)), dim=1)  # type: ignore[arg-type]
+            out = self.cat.cat([x1, self.branch2(x2)], dim=1)
         else:
-            out = self.cat.cat((self.branch1(x), self.branch2(x)), dim=1)  # type: ignore[arg-type]
+            out = self.cat.cat([self.branch1(x), self.branch2(x)], dim=1)
 
         out = shufflenetv2.channel_shuffle(out, 2)
 
@@ -40,7 +40,12 @@ def forward(self, x: Tensor) -> Tensor:
 
 class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super(QuantizableShuffleNetV2, self).__init__(*args, inverted_residual=QuantizableInvertedResidual, **kwargs)
+        # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
+        super(QuantizableShuffleNetV2, self).__init__(  # type: ignore[misc]
+            *args,
+            inverted_residual=QuantizableInvertedResidual,
+            **kwargs
+        )
         self.quant = torch.quantization.QuantStub()
         self.dequant = torch.quantization.DeQuantStub()
 

From a6b852882a8714afbaf2c73001543c32dc0dcd89 Mon Sep 17 00:00:00 2001
From: Philip Meier <github.pmeier@posteo.de>
Date: Wed, 18 Aug 2021 07:23:58 +0200
Subject: [PATCH 13/13] cleanup #TODO comments

---
 torchvision/models/quantization/inception.py    | 8 ++++++--
 torchvision/models/quantization/mobilenetv3.py  | 2 +-
 torchvision/models/quantization/shufflenetv2.py | 2 +-
 torchvision/models/quantization/utils.py        | 1 +
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
index 71bcfc7b4fd..6c6384c295a 100644
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -106,10 +106,9 @@ def forward(self, x: Tensor) -> Tensor:
     def fuse_model(self) -> None:
         torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True)
 
-# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
-
 
 class QuantizableInceptionA(inception_module.InceptionA):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionA, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
@@ -124,6 +123,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionB(inception_module.InceptionB):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionB, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
@@ -138,6 +138,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionC(inception_module.InceptionC):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionC, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
@@ -152,6 +153,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionD(inception_module.InceptionD):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionD, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
@@ -166,6 +168,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionE(inception_module.InceptionE):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionE, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
@@ -203,6 +206,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableInceptionAux(inception_module.InceptionAux):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(QuantizableInceptionAux, self).__init__(  # type: ignore[misc]
             conv_block=QuantizableBasicConv2d,
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
index 050db77001e..a1aa9d7d4bd 100644
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/torchvision/models/quantization/mobilenetv3.py
@@ -28,8 +28,8 @@ def fuse_model(self) -> None:
         fuse_modules(self, ['fc1', 'relu'], inplace=True)
 
 
-# TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
 class QuantizableInvertedResidual(InvertedResidual):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(  # type: ignore[misc]
             se_layer=QuantizableSqueezeExcitation,
diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py
index 16be4f21ce5..4f0861dcb30 100644
--- a/torchvision/models/quantization/shufflenetv2.py
+++ b/torchvision/models/quantization/shufflenetv2.py
@@ -39,8 +39,8 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class QuantizableShuffleNetV2(shufflenetv2.ShuffleNetV2):
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
         super(QuantizableShuffleNetV2, self).__init__(  # type: ignore[misc]
             *args,
             inverted_residual=QuantizableInvertedResidual,
diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py
index 941db439e0c..c195d162482 100644
--- a/torchvision/models/quantization/utils.py
+++ b/torchvision/models/quantization/utils.py
@@ -32,6 +32,7 @@ def quantize_model(model: nn.Module, backend: str) -> None:
             activation=torch.quantization.default_observer,
             weight=torch.quantization.default_weight_observer)
 
+    # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659
     model.fuse_model()  # type: ignore[operator]
     torch.quantization.prepare(model, inplace=True)
     model(_dummy_input_data)