(self, *inputs, **kwargs)
| 495 | self.module.float() |
| 496 | |
| 497 | def _create_cuda_graph(self, *inputs, **kwargs): |
| 498 | # warmup to create the workspace and cublas handle |
| 499 | cuda_stream = get_accelerator().Stream() |
| 500 | cuda_stream.wait_stream(get_accelerator().current_stream()) |
| 501 | with get_accelerator().stream(cuda_stream): |
| 502 | for i in range(3): |
| 503 | ret = self.module(*inputs, **kwargs) |
| 504 | get_accelerator().current_stream().wait_stream(cuda_stream) |
| 505 | |
| 506 | # create cuda_graph and assign static_inputs and static_outputs |
| 507 | self._cuda_graphs = get_accelerator().create_graph() |
| 508 | self.static_inputs = inputs |
| 509 | self.static_kwargs = kwargs |
| 510 | |
| 511 | with get_accelerator().capture_to_graph(self._cuda_graphs): |
| 512 | self.static_output = self.module(*self.static_inputs, **self.static_kwargs) |
| 513 | |
| 514 | self.cuda_graph_created = True |
| 515 | |
| 516 | def _graph_replay(self, *inputs, **kwargs): |
| 517 | for i in range(len(inputs)): |
no test coverage detected