From dba06b36a1970b891d8dcfe086ad5d57715392dc Mon Sep 17 00:00:00 2001 From: "Wu, Zhenyu" Date: Thu, 6 Aug 2026 17:15:33 +0800 Subject: [PATCH] fix: fix wrong dtype --- compressai/models/google.py | 2 +- compressai/models/vbr.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/compressai/models/google.py b/compressai/models/google.py index 712ce031..6e6a2de2 100644 --- a/compressai/models/google.py +++ b/compressai/models/google.py @@ -724,7 +724,7 @@ def _decompress_ar( rv = decoder.decode_stream( indexes.squeeze().tolist(), cdf, cdf_lengths, offsets ) - rv = torch.Tensor(rv).reshape(1, -1, 1, 1) + rv = torch.tensor(rv, dtype=means_hat.dtype, device=means_hat.device).reshape(1, -1, 1, 1) rv = self.gaussian_conditional.dequantize(rv, means_hat) hp = h + padding diff --git a/compressai/models/vbr.py b/compressai/models/vbr.py index 35769333..8fb250a1 100644 --- a/compressai/models/vbr.py +++ b/compressai/models/vbr.py @@ -921,9 +921,8 @@ def _decompress_ar( # noqa: C901 rv = decoder.decode_stream( indexes.squeeze().tolist(), cdf, cdf_lengths, offsets ) - rv = ( - torch.Tensor(rv).reshape(1, -1, 1, 1).to(scales_hat.device) - ) # TODO: move rv to gpu ? + # TODO: move rv to gpu ? + rv = torch.tensor(rv, dtype=means_hat.dtype, device=means_hat.device).reshape(1, -1, 1, 1) if stage == 1: rv = self.gaussian_conditional.dequantize(rv, means_hat)