(tensor: Tensor)
| 33 | |
| 34 | |
| 35 | def tensor_round(tensor: Tensor) -> Tensor: |
| 36 | if tensor.device.type == 'cpu' and tensor.dtype == torch.float16: |
| 37 | # CPU does not support FP16 round |
| 38 | return tensor.to(dtype=torch.float32).round_().to(dtype=torch.float16) |
| 39 | else: |
| 40 | return tensor.round_() |
| 41 | |
| 42 | |
| 43 | class Quantizer: |