[quant][bug] Fix histogram observer with 0 input (pytorch#40191)

Summary:
Pull Request resolved: pytorch#40191

When the first couple of inputs passed to histogram observer are all 0's subsequent non-zero inputs cause a div by 0 error

Test Plan:
python test/test_quantization.py TestHistogramObserver.test_histogram_observer_zero_inputs

Imported from OSS

Differential Revision: D22119422

fbshipit-source-id: 8bbbba914ba7f343121830c768ca0444439f8e03

Author: supriyar

test/quantization/test_workflow_module.py

@@ -397,6 +397,20 @@ def test_histogram_observer_one_sided(self):
         qparams = myobs.calculate_qparams()
         self.assertEqual(qparams[1].item(), 0)
 
+    def test_histogram_observer_zero_inputs(self):
+        myobs = HistogramObserver(bins=3, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, reduce_range=False)
+        x = torch.zeros(4, requires_grad=True)
+        y = torch.tensor([2.0, 3.0, 4.0, 5.0], requires_grad=True)
+        z = torch.tensor([5.0, 6.0, 7.0, 8.0])
+        myobs(x)
+        myobs(x)
+        myobs(y)
+        myobs(z)
+        qparams = myobs.calculate_qparams()
+        self.assertEqual(myobs.min_val, 2.0)
+        self.assertEqual(myobs.max_val, 8.0)
+        self.assertEqual(myobs.histogram, [2., 3., 3.])
+
 class TestFakeQuantizePerTensor(TestCase):
     @given(device=st.sampled_from(['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu']),
            X=hu.tensor(shapes=hu.array_shapes(1, 5,),

torch/quantization/observer.py

@@ -828,7 +828,10 @@ def forward(self, x_orig):
         x = x_orig.detach()
         min_val = self.min_val
         max_val = self.max_val
-        if min_val.numel() == 0 or max_val.numel() == 0:
+        prev_zeros = False
+        if min_val.numel() > 0 and max_val.numel() > 0:
+            prev_zeros = (min_val.item() == 0) and (max_val.item() == 0)
+        if min_val.numel() == 0 or max_val.numel() == 0 or prev_zeros:
             min_val = torch.min(x)
             max_val = torch.max(x)
             self.min_val.resize_(min_val.shape)