Update tasks/model.py

tasks/model.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 from torch.utils.data import DataLoader, TensorDataset
 from torchaudio import transforms
 from torchvision import models
+from torch.quantization import QuantStub, DeQuantStub
 
 
 
@@ -104,4 +105,19 @@ class BlazeFaceModel(nn.Module):
     def forward(self, x):
         features = self.blazeface_backbone(x)
         output = self.fc(features)
-        return output
+        return output
+
+class QuantizedBlazeFaceModel(nn.Module):
+        def __init__(self, model_fp32):
+            super(QuantizedBlazeFaceModel, self).__init__()
+            self.quant = QuantStub()
+            self.dequant = DeQuantStub()
+            self.backbone = model_fp32.blazeface_backbone
+            self.fc = model_fp32.fc
+    
+        def forward(self, x):
+            x = self.quant(x)
+            x = self.backbone(x)
+            x = self.fc(x)
+            x = self.dequant(x)
+            return x