fix tensorrt int8 (#1799)

open-mmlab · Feb 22, 2023 · 7b45556 · 7b45556
1 parent 0b9a9e2
commit 7b45556
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mmdeploy/backend/tensorrt/utils.py b/mmdeploy/backend/tensorrt/utils.py
@@ -140,11 +140,13 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto],
         >>>             })
     """
 
+    if int8_mode or device_id != 0:
+        import pycuda.autoinit  # noqa:F401
+
     if device_id != 0:
         import os
         old_cuda_device = os.environ.get('CUDA_DEVICE', None)
         os.environ['CUDA_DEVICE'] = str(device_id)
-        import pycuda.autoinit  # noqa:F401
         if old_cuda_device is not None:
             os.environ['CUDA_DEVICE'] = old_cuda_device
         else: