Skip to content

Commit 6865d19

Browse files
authored
TensorRT 7 export fix (#6235)
1 parent 33a67b4 commit 6865d19

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

export.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
174174
check_requirements(('tensorrt',))
175175
import tensorrt as trt
176176

177-
if trt.__version__[0] == 7: # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
177+
if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
178178
grid = model.model[-1].anchor_grid
179179
model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
180180
export_onnx(model, im, file, 12, train, False, simplify) # opset 12

0 commit comments

Comments
 (0)