diff --git a/cdvae/pl_data/datamodule.py b/cdvae/pl_data/datamodule.py
index cd9a4380..36f296d5 100644
--- a/cdvae/pl_data/datamodule.py
+++ b/cdvae/pl_data/datamodule.py
@@ -94,7 +94,7 @@ def setup(self, stage: Optional[str] = None):
                 for dataset_cfg in self.datasets.test
             ]
             for test_dataset in self.test_datasets:
-                test_dataset.lattice_scaler = self.scaler
+                test_dataset.lattice_scaler = self.lattice_scaler
                 test_dataset.scaler = self.scaler
 
     def train_dataloader(self) -> DataLoader: