| 290 | self.context = self.engine.create_execution_context() |
| 291 | |
| 292 | def allocate_buffers(self, shape_dict=None, device="cuda"): |
| 293 | nvtx.range_push("allocate_buffers") |
| 294 | for idx in range(self.engine.num_io_tensors): |
| 295 | binding = self.engine[idx] |
| 296 | if shape_dict and binding in shape_dict: |
| 297 | shape = shape_dict[binding].shape |
| 298 | else: |
| 299 | shape = self.context.get_binding_shape(idx) |
| 300 | dtype = trt.nptype(self.engine.get_binding_dtype(binding)) |
| 301 | if self.engine.binding_is_input(binding): |
| 302 | self.context.set_binding_shape(idx, shape) |
| 303 | tensor = torch.empty( |
| 304 | tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype] |
| 305 | ).to(device=device) |
| 306 | self.tensors[binding] = tensor |
| 307 | nvtx.range_pop() |
| 308 | |
| 309 | def infer(self, feed_dict, stream, use_cuda_graph=False): |
| 310 | nvtx.range_push("set_tensors") |