(weights)
| 103 | |
| 104 | |
| 105 | def build_engine(weights): |
| 106 | # For more information on TRT basics, refer to the introductory samples. |
| 107 | builder = trt.Builder(TRT_LOGGER) |
| 108 | network = builder.create_network(common.EXPLICIT_BATCH) |
| 109 | config = builder.create_builder_config() |
| 110 | runtime = trt.Runtime(TRT_LOGGER) |
| 111 | |
| 112 | config.max_workspace_size = common.GiB(1) |
| 113 | # Populate the network using weights from the PyTorch model. |
| 114 | populate_network(network, weights) |
| 115 | # Build and return an engine. |
| 116 | plan = builder.build_serialized_network(network, config) |
| 117 | return runtime.deserialize_cuda_engine(plan) |
| 118 | |
| 119 | |
| 120 | # Loads a random test case from pytorch's DataLoader |
no test coverage detected