diff --git a/lato_mlx/models/vvae.py b/lato_mlx/models/vvae.py index 5065786..02f2335 100644 --- a/lato_mlx/models/vvae.py +++ b/lato_mlx/models/vvae.py @@ -46,7 +46,10 @@ class DownResBlock(nn.Module): self.block = SparseResBlock(channels, out_channels) def __call__(self, x: SparseTensor) -> SparseTensor: - return self.block(downsample(x, 2)) + # mode="max" explicitly: LATO.2 hardcodes reduce="amax", but Pixal3D's version of + # the same op defaults to "mean". Relying on the shared default would silently + # change every downsampled feature the day that default moves. + return self.block(downsample(x, 2, mode="max")) class SparseTransformerBase(nn.Module):