| 104 | self.prealloc_buffers(self.output_names) |
| 105 | |
| 106 | def run_async(self, stream=None, **inputs): |
| 107 | if stream is None: |
| 108 | stream = torch.cuda.current_stream() |
| 109 | |
| 110 | for name, tensor in inputs.items(): |
| 111 | self.mgx_args[name] = self.tensor_to_arg(tensor) |
| 112 | |
| 113 | self.start_events.append(torch.cuda.Event(enable_timing=True)) |
| 114 | self.end_events.append(torch.cuda.Event(enable_timing=True)) |
| 115 | |
| 116 | self.start_events[-1].record() |
| 117 | self.model.run_async(self.mgx_args, stream.cuda_stream, "ihipStream_t") |
| 118 | self.end_events[-1].record() |
| 119 | |
| 120 | return {p: self.torch_buffers[p] for p in self.output_names} |
| 121 | |
| 122 | def save_model(self, path): |
| 123 | os.makedirs(os.path.dirname(path), exist_ok=True) |