(self, x)
| 113 | |
| 114 | @T.no_grad() |
| 115 | def quantize(self, x): |
| 116 | if self.c.split: |
| 117 | x1, x2 = x.chunk(2, dim=-1) |
| 118 | with T.autocast(device_type='cuda', dtype=T.bfloat16): |
| 119 | quantized1 = self.quantizer(x1) |
| 120 | quantized2 = self.quantizer(x2) |
| 121 | return quantized1, quantized2 |
| 122 | else: |
| 123 | with T.autocast(device_type='cuda', dtype=T.bfloat16): |
| 124 | return self.quantizer(x) |
| 125 | |
| 126 | @T.no_grad() |
| 127 | def untokenize(self, token_data): |