gpu support

Author: tokkiwa

compressai/entropy_models/entropy_models.py

@@ -878,10 +878,10 @@ def compress(self, y, scales, means, weights):
         #     abs_max + 1,
         # )
         rv = self.entropy_coder._encoder.encode_with_indexes_gmm(
-            symbols,
-            scales, 
-            means,
-            weights,
+            symbols.to("cpu"),
+            scales.to("cpu"), 
+            means.to("cpu"),
+            weights.to("cpu"),
             abs_max + 1
         )
 
@@ -900,9 +900,9 @@ def decompress(self, strings, abs_max, zero_bitmap, scales, means, weights):
         # )
         symbols = self.entropy_coder._decoder.decode_with_indexes_gmm(
             strings,
-            scales,
-            means,
-            weights,
+            scales.to("cpu"),
+            means.to("cpu"),
+            weights.to("cpu"),
             abs_max + 1
         )