diff --git a/tests/utils.py b/tests/utils.py index 4bda8955ed..cf1711292f 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -724,22 +724,16 @@ def test_script_save(net, *inputs, device=None, rtol=1e-4, atol=0.0): """ # TODO: would be nice to use GPU if available, but it currently causes CI failures. device = "cpu" - try: - with tempfile.TemporaryDirectory() as tempdir: - convert_to_torchscript( - model=net, - filename_or_obj=os.path.join(tempdir, "model.ts"), - verify=True, - inputs=inputs, - device=device, - rtol=rtol, - atol=atol, - ) - except (RuntimeError, AttributeError): - if sys.version_info.major == 3 and sys.version_info.minor == 11: - warnings.warn("skipping py 3.11") - return - raise + with tempfile.TemporaryDirectory() as tempdir: + convert_to_torchscript( + model=net, + filename_or_obj=os.path.join(tempdir, "model.ts"), + verify=True, + inputs=inputs, + device=device, + rtol=rtol, + atol=atol, + ) def test_onnx_save(net, *inputs, device=None, rtol=1e-4, atol=0.0): @@ -753,23 +747,17 @@ def test_onnx_save(net, *inputs, device=None, rtol=1e-4, atol=0.0): # TODO: would be nice to use GPU if available, but it currently causes CI failures. device = "cpu" _, has_onnxruntime = optional_import("onnxruntime") - try: - with tempfile.TemporaryDirectory() as tempdir: - convert_to_onnx( - model=net, - filename=os.path.join(tempdir, "model.onnx"), - verify=True, - inputs=inputs, - device=device, - use_ort=has_onnxruntime, - rtol=rtol, - atol=atol, - ) - except (RuntimeError, AttributeError): - if sys.version_info.major == 3 and sys.version_info.minor == 11: - warnings.warn("skipping py 3.11") - return - raise + with tempfile.TemporaryDirectory() as tempdir: + convert_to_onnx( + model=net, + filename=os.path.join(tempdir, "model.onnx"), + verify=True, + inputs=inputs, + device=device, + use_ort=has_onnxruntime, + rtol=rtol, + atol=atol, + ) def download_url_or_skip_test(*args, **kwargs):