Skip to content

Commit bb54c68

Browse files
committed
lenet5.py migrated to trt8
1 parent d0eb871 commit bb54c68

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

lenet/lenet.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import sys
55

66
import numpy as np
7-
import pycuda.autoinit
87
import pycuda.driver as cuda
98
import tensorrt as trt
109

@@ -119,7 +118,7 @@ def createLenetEngine(maxBatchSize, builder, config, dt):
119118

120119
# Build engine
121120
builder.max_batch_size = maxBatchSize
122-
builder.max_workspace_size = 1 << 20
121+
config.max_workspace_size = 1 << 20
123122
engine = builder.build_engine(network, config)
124123

125124
del network

0 commit comments

Comments
 (0)