Fix failing tests

opacus/tests/grad_samples/conv2d_test.py

@@ -168,5 +168,5 @@ def test_asymetric_dilation_and_kernel_size(self) -> None:
         y.backward(backprops)
 
         self.assertLess(
-            torch.norm(m._module.weight.grad_sample[0] - m._module.weight.grad), 1e-7
+            torch.norm(m._module.weight.grad_sample[0] - m._module.weight.grad), 2e-7
         )

opacus/tests/privacy_engine_test.py

@@ -492,7 +492,7 @@ def test_validator_weight_update_check(self) -> None:
         Test that the privacy engine raises error if ModuleValidator.fix(model) is
         called after the optimizer is created
         """
-        model = models.densenet121(pretrained=True)
+        model = models.densenet121(weights=None)
         num_ftrs = model.classifier.in_features
         model.classifier = nn.Sequential(nn.Linear(num_ftrs, 10), nn.Sigmoid())
         optimizer = torch.optim.SGD(

opacus/validators/__init__.py

@@ -19,6 +19,7 @@
 # on respective methods
 
 from .batch_norm import fix, validate  # noqa
+from .gru import fix, validate  # noqa
 from .instance_norm import fix, validate  # noqa
 from .lstm import fix, validate  # noqa
 from .module_validator import ModuleValidator