CheckerboardLatentCodec broken in torch<2.0 about compressai HOT 2 OPEN

Thanks for the report.

I could add a version check for torch<2.0:

from packaging.version import Version

class CheckerboardLatentCodec(LatentCodec):
    def _y_ctx_zero(self, y: Tensor) -> Tensor:
        if Version(torch.__version__) < Version("2.0.0"):
            return self._mask(self.context_prediction(y).detach(), "all")
        return y.new_zeros(self.context_prediction(y.to("meta")).shape)

...but perhaps simpler is just to revert:

class CheckerboardLatentCodec(LatentCodec):
    def _y_ctx_zero(self, y: Tensor) -> Tensor:
        return self._mask(self.context_prediction(y).detach(), "all")

from compressai.

To be fair I don't actually know which is the earliest torch version that supports meta device tensors as I couldn't find any solid information.

Although I think the simpler fix is probably good enough. On my machine with a 14900K and a 3090 and an (unreasonably large) context size of (16, 192, 512, 512) it takes 0.06ms to execute that line on GPU. It does take about 4 seconds on CPU, but with a more reasonable context size of (16, 192, 32, 32) it takes roughly 80ms on CPU.

from compressai.