Allocates outputs to run TRT engine Args: device: GPU device to allocate memory on
(self, device)
| 151 | ) |
| 152 | |
| 153 | def allocate_buffers(self, device): |
| 154 | """ |
| 155 | Allocates outputs to run TRT engine |
| 156 | Args: |
| 157 | device: GPU device to allocate memory on |
| 158 | """ |
| 159 | ctx = self.context |
| 160 | |
| 161 | for i, binding in enumerate(self.output_names): |
| 162 | shape = list(ctx.get_tensor_shape(binding)) |
| 163 | if binding not in self.tensors or list(self.tensors[binding].shape) != shape: |
| 164 | t = torch.empty(shape, dtype=self.dtypes[i], device=device).contiguous() |
| 165 | self.tensors[binding] = t |
| 166 | ctx.set_tensor_address(binding, t.data_ptr()) |
| 167 | |
| 168 | def set_inputs(self, feed_dict, stream): |
| 169 | """ |