Added number of features in FrozenBatchNorm2d __repr__

test/test_ops.py

@@ -538,5 +538,15 @@ def script_func(x_, offset_, weight_, bias_, stride_, pad_, dilation_):
                   (x, offset, weight, bias), nondet_tol=1e-5)
 
 
+class FrozenBNTester(unittest.TestCase):
+    def test_frozenbatchnorm2d_repr(self):
+        num_features = 32
+        t = ops.misc.FrozenBatchNorm2d(num_features)
+
+        # Check integrity of object __repr__ attribute
+        expected_string = 'FrozenBatchNorm2d({0})'.format(num_features)
+        self.assertEqual(t.__repr__(), expected_string)
+
+
 if __name__ == '__main__':
     unittest.main()

torchvision/ops/misc.py

@@ -151,3 +151,6 @@ def forward(self, x):
         scale = w * rv.rsqrt()
         bias = b - rm * scale
         return x * scale + bias
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.weight.shape[0]})"