(tensor: torch.Tensor, num_beams: int)
| 168 | |
| 169 | |
| 170 | def _tile_beam_width(tensor: torch.Tensor, num_beams: int): |
| 171 | new_shape = np.array(tensor.shape) |
| 172 | new_shape[0] = new_shape[0] * num_beams |
| 173 | |
| 174 | tile_size = np.ones(new_shape.shape, dtype=np.int32) |
| 175 | tile_size = np.insert(tile_size, 1, num_beams) |
| 176 | |
| 177 | new_tensor = torch.unsqueeze(tensor, 1) |
| 178 | new_tensor = new_tensor.tile(tile_size.tolist()) |
| 179 | new_tensor = new_tensor.reshape(new_shape.tolist()) |
| 180 | return new_tensor |
| 181 | |
| 182 | |
| 183 | class _Profiler(trt.IProfiler): |
no test coverage detected