| 307 | nvtx.range_pop() |
| 308 | |
| 309 | def infer(self, feed_dict, stream, use_cuda_graph=False): |
| 310 | nvtx.range_push("set_tensors") |
| 311 | for name, buf in feed_dict.items(): |
| 312 | self.tensors[name].copy_(buf) |
| 313 | |
| 314 | for name, tensor in self.tensors.items(): |
| 315 | self.context.set_tensor_address(name, tensor.data_ptr()) |
| 316 | nvtx.range_pop() |
| 317 | nvtx.range_push("execute") |
| 318 | noerror = self.context.execute_async_v3(stream) |
| 319 | if not noerror: |
| 320 | raise ValueError("ERROR: inference failed.") |
| 321 | nvtx.range_pop() |
| 322 | return self.tensors |
| 323 | |
| 324 | def __str__(self): |
| 325 | out = "" |