diff --git a/geofree/main.py b/geofree/main.py
index 80e8b77..f5938a7 100644
--- a/geofree/main.py
+++ b/geofree/main.py
@@ -9,7 +9,7 @@
 from pytorch_lightning import seed_everything
 from pytorch_lightning.trainer import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
 
 
 def get_obj_from_str(string, reload=False):