diff --git a/compressai/models/google.py b/compressai/models/google.py index 712ce031..f4b60b83 100644 --- a/compressai/models/google.py +++ b/compressai/models/google.py @@ -633,7 +633,7 @@ def _compress_ar(self, y_hat, params, height, width, kernel_size, padding): y_crop = y_crop[:, :, padding, padding] y_q = self.gaussian_conditional.quantize(y_crop, "symbols", means_hat) - y_hat[:, :, h + padding, w + padding] = y_q + means_hat + y_hat[:, :, h + padding, w + padding] = self.gaussian_conditional.dequantize(y_q, means_hat) symbols_list.extend(y_q.squeeze().tolist()) indexes_list.extend(indexes.squeeze().tolist())