diff --git a/cldm/logger.py b/cldm/logger.py
index 385eb31..bf5ae6c 100644
--- a/cldm/logger.py
+++ b/cldm/logger.py
@@ -15,7 +15,7 @@
 import torchvision
 from PIL import Image
 from pytorch_lightning.callbacks import Callback
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
 import pytorch_lightning as pl
 
 class ImageLogger(Callback):
diff --git a/ldm/models/diffusion/ddpm.py b/ldm/models/diffusion/ddpm.py
index 62dc7ea..7ce62a8 100644
--- a/ldm/models/diffusion/ddpm.py
+++ b/ldm/models/diffusion/ddpm.py
@@ -27,7 +27,7 @@
 import itertools
 from tqdm import tqdm
 from torchvision.utils import make_grid
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
 from omegaconf import ListConfig
 
 from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
diff --git a/ldm/models/diffusion/ddpm_multi.py b/ldm/models/diffusion/ddpm_multi.py
index 516a65f..5bdcb15 100644
--- a/ldm/models/diffusion/ddpm_multi.py
+++ b/ldm/models/diffusion/ddpm_multi.py
@@ -27,7 +27,7 @@
 import itertools
 from tqdm import tqdm
 from torchvision.utils import make_grid
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
 from omegaconf import ListConfig
 
 from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config