glenn-jocher commited on
Commit
1f1ec1c
·
unverified ·
1 Parent(s): 404b4fe

Restrict TRT autoinstall to Linux-only (#7549)

Browse files

May partially resolve concerns in https://github.com/ultralytics/yolov5/pull/7537#discussion_r856843711

Files changed (1) hide show
  1. export.py +2 -1
export.py CHANGED
@@ -218,7 +218,8 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
218
  # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
219
  try:
220
  assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
221
- check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
 
222
  import tensorrt as trt
223
 
224
  if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
 
218
  # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
219
  try:
220
  assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
221
+ if platform.system() == 'Linux':
222
+ check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
223
  import tensorrt as trt
224
 
225
  if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012