(self, a, a_compressed, meta, run_on_cpu: bool)
| 78 | self.__init__(a.shape[0], a.shape[1], a.shape[2]) |
| 79 | |
| 80 | def compress(self, a, a_compressed, meta, run_on_cpu: bool): |
| 81 | if run_on_cpu: |
| 82 | if a.device.type != "cpu": |
| 83 | raise ValueError("a must be on cpu") |
| 84 | return self.__compress_on_cpu(a, a_compressed, meta) |
| 85 | else: |
| 86 | if a.device.type != "cuda": |
| 87 | raise ValueError("a must be on cuda") |
| 88 | return self.__compress_on_cuda(a, a_compressed, meta) |
| 89 | |
| 90 | def __compress_on_cpu(self, a, a_compressed, meta): |
| 91 | """ |