diff --git a/tests/test_apis/test_onnx2tensorrt.py b/tests/test_apis/test_onnx2tensorrt.py index c84399530..b47ae0ce8 100644 --- a/tests/test_apis/test_onnx2tensorrt.py +++ b/tests/test_apis/test_onnx2tensorrt.py @@ -34,7 +34,7 @@ def get_deploy_cfg(): backend_config=dict( type='tensorrt', common_config=dict( - fp16_mode=False, max_workspace_size=1 << 30), + fp16_mode=False, max_workspace_size=1 << 20), model_inputs=[ dict( input_shapes=dict( diff --git a/tests/test_codebase/test_mmocr/test_mmocr_models.py b/tests/test_codebase/test_mmocr/test_mmocr_models.py index 16ffa0dca..04d916943 100644 --- a/tests/test_codebase/test_mmocr/test_mmocr_models.py +++ b/tests/test_codebase/test_mmocr/test_mmocr_models.py @@ -403,7 +403,7 @@ def test_forward_of_fpnc(backend: Backend): dict( backend_config=dict( type=backend.value, - common_config=dict(max_workspace_size=1 << 30), + common_config=dict(max_workspace_size=1 << 20), model_inputs=[ dict( input_shapes=dict( @@ -550,7 +550,7 @@ def test_mmdet_wrapper__forward(backend): dict( backend_config=dict( type=backend.value, - common_config=dict(max_workspace_size=1 << 30)), + common_config=dict(max_workspace_size=1 << 20)), onnx_config=dict( input_shape=None, input_names=['inputs'], diff --git a/tests/test_ops/utils.py b/tests/test_ops/utils.py index c36f678e0..0291158e1 100644 --- a/tests/test_ops/utils.py +++ b/tests/test_ops/utils.py @@ -119,7 +119,7 @@ class TestTensorRTExporter: backend_config=dict( type='tensorrt', common_config=dict( - fp16_mode=False, max_workspace_size=1 << 28), + fp16_mode=False, max_workspace_size=1 << 20), model_inputs=[ dict( input_shapes=dict( @@ -132,14 +132,13 @@ class TestTensorRTExporter: ]))) ]))) - onnx_model = onnx.load(onnx_file_path) work_dir, filename = os.path.split(trt_file_path) trt_apis.onnx2tensorrt( work_dir, filename, 0, deploy_cfg=deploy_cfg, - onnx_model=onnx_model) + onnx_model=onnx_file_path) if expected_result is None and not isinstance( model, onnx.onnx_ml_pb2.ModelProto): with torch.no_grad():