feat: Add QAT patch which modifies scale factor dtype to INT32

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

docker/Dockerfile

@@ -26,15 +26,15 @@ COPY . /workspace/trtorch/src
 WORKDIR /workspace/trtorch/src
 RUN cp ./docker/WORKSPACE.cu.docker  WORKSPACE
 
-# This script builds both libtrtorch bin/lib/include tarball and the Pythin wheel, in dist/  
+# This script builds both libtrtorch bin/lib/include tarball and the Pythin wheel, in dist/
 RUN ./docker/dist-build.sh
 
 FROM base as trtorch
 
 # copy source repo
 COPY . /workspace/trtorch
 COPY --from=trtorch-builder  /workspace/trtorch/src/dist/ .
-
+RUN patch -u /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py -i /workspace/trtorch/docker/qat.patch
 RUN conda init bash
 
 RUN pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org

docker/qat.patch

@@ -0,0 +1,11 @@
+--- /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py	2021-08-16 22:50:37.000000000 +0000
++++ tensor_quantizer.py	2021-10-19 20:41:54.288077426 +0000
+@@ -291,7 +291,7 @@
+             quant_dim = list(amax.shape).index(list(amax_sequeeze.shape)[0])
+             scale = amax_sequeeze / bound
+             outputs = torch.fake_quantize_per_channel_affine(
+-                inputs, scale.data, torch.zeros_like(scale, dtype=torch.long).data, quant_dim,
++                inputs, scale.data, torch.zeros_like(scale, dtype=torch.int32).data, quant_dim,
+                 -bound - 1 if not self._unsigned else 0, bound)
+ 
+         return outputs