(self, input_desc, output_desc, inputs, outputs, workspace, stream)
| 397 | ) |
| 398 | |
| 399 | def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream): |
| 400 | torch_stream = torch.cuda.ExternalStream(stream_ptr=stream) |
| 401 | self.workspace = workspace |
| 402 | self.current_stream = stream |
| 403 | |
| 404 | with torch.cuda.stream(torch_stream): |
| 405 | self.forward( |
| 406 | tuple( |
| 407 | TensorWrapper.from_trt_desc(input_desc[i], inputs[i]) |
| 408 | for i in range(len(input_desc)) |
| 409 | ), |
| 410 | tuple( |
| 411 | TensorWrapper.from_trt_desc(output_desc[i], outputs[i]) |
| 412 | for i in range(len(output_desc)) |
| 413 | ), |
| 414 | ) |
| 415 | |
| 416 | self.current_stream = -1 |
| 417 | |
| 418 | def __call__(self, *args: Union[Sequence[TensorWrapper], Sequence[torch.Tensor]]): |
| 419 | is_trtllm = True |
nothing calls this directly
no test coverage detected