(tensor: Tensor, min, max)
| 25 | |
| 26 | |
| 27 | def tensor_clamp(tensor: Tensor, min, max) -> Tensor: |
| 28 | if tensor.device.type == 'cpu' and tensor.dtype == torch.float16: |
| 29 | # CPU does not support FP16 clamp |
| 30 | return tensor.to(dtype=torch.float32).clamp_(min, max).to(dtype=torch.float16) |
| 31 | else: |
| 32 | return tensor.clamp_(min, max) |
| 33 | |
| 34 | |
| 35 | def tensor_round(tensor: Tensor) -> Tensor: |