fix: fix wrong dtype

Freed-Wu · fracape · commit 6cc371fde518 · 2026-03-12T15:57:25.000-07:00
diff --git a/compressai/models/google.py b/compressai/models/google.py
@@ -673,6 +673,7 @@ def decompress(self, strings, shape):
         y_hat = torch.zeros(
             (z_hat.size(0), self.M, y_height + 2 * padding, y_width + 2 * padding),
             device=z_hat.device,
+            dtype=z_hat.dtype,
         )
 
         for i, y_string in enumerate(strings[0]):

Original file line number	Diff line number	Diff line change
`@@ -673,6 +673,7 @@ def decompress(self, strings, shape):`
`673`	`673`	`y_hat = torch.zeros(`
`674`	`674`	`(z_hat.size(0), self.M, y_height + 2 * padding, y_width + 2 * padding),`
`675`	`675`	`device=z_hat.device,`
	`676`	`+ dtype=z_hat.dtype,`
`676`	`677`	`)`
`677`	`678`
`678`	`679`	`for i, y_string in enumerate(strings[0]):`