diff --git a/tests/pytorch_tests/model_tests/feature_models/bn_folding_test.py b/tests/pytorch_tests/model_tests/feature_models/bn_folding_test.py index bf3edb67e..ed217a5c7 100644 --- a/tests/pytorch_tests/model_tests/feature_models/bn_folding_test.py +++ b/tests/pytorch_tests/model_tests/feature_models/bn_folding_test.py @@ -85,12 +85,7 @@ def compare(self, quantized_models, float_model, input_x=None, quantization_info is_bn_in_model = nn.BatchNorm2d in [type(module) for name, module in quant_model.named_modules()] self.unit_test.assertTrue(self.fold_applied is not is_bn_in_model) - - # TODO: remove this and set atol=1e-5 for all tests after fixing the issue with ConvTranspose2d. - # https://github.com/sony/model_optimization/issues/1115 - is_convtranspose2d_in_model = nn.ConvTranspose2d in [type(module) for name, module in quant_model.named_modules()] - atol = 1e-3 if is_convtranspose2d_in_model else 1e-5 - self.unit_test.assertTrue(np.isclose(out_quant, out_float, atol=atol, rtol=1e-4).all()) + self.unit_test.assertTrue(np.isclose(out_quant, out_float, atol=1e-5, rtol=1e-4).all()) class BNForwardFoldingNet(nn.Module):