diff --git a/examples/images/diffusion/README.md b/examples/images/diffusion/README.md
index cc57f6d54..9e9900f2c 100644
--- a/examples/images/diffusion/README.md
+++ b/examples/images/diffusion/README.md
@@ -69,7 +69,7 @@ python setup.py install
 ##### From pip
 
 ```
-pip install pytorch-lightning
+pip install pytorch-lightning==2.0.0rc0
 ```
 
 #### Step 3:Install [Colossal-AI](https://colossalai.org/download/) From Our Official Website
diff --git a/examples/images/diffusion/environment.yaml b/examples/images/diffusion/environment.yaml
index 5164be72e..d1ec69c1a 100644
--- a/examples/images/diffusion/environment.yaml
+++ b/examples/images/diffusion/environment.yaml
@@ -27,5 +27,6 @@ dependencies:
     - torchmetrics==0.7.0
     - prefetch_generator
     - datasets
-    - colossalai
+    - colossalai==0.2.5
+    - lightning==1.9.0
     - -e .