diff --git a/benchmarks/lightning/main.py b/benchmarks/lightning/main.py
index 9295c504..97aae5c9 100644
--- a/benchmarks/lightning/main.py
+++ b/benchmarks/lightning/main.py
@@ -1,10 +1,13 @@
 #!/usr/bin/env python
 
-from habana_frameworks.torch import hpu; hpu.init()
 
 import argparse
 import os
 
+os.environ["PT_HPU_LAZY_MODE"] = str(int(os.getenv("WORLD_SIZE", -1) <= 0))
+
+from habana_frameworks.torch import hpu; hpu.init()
+
 import torch
 import torch.nn.functional as F
 import lightning as L
diff --git a/docker/Makefile b/docker/Makefile
index a63fcb91..6826d1b3 100644
--- a/docker/Makefile
+++ b/docker/Makefile
@@ -13,5 +13,5 @@ hpu:
 	# docker system prune -a -f
 	# docker image prune -a -f
 	docker build --build-arg BENCH=$(bench) --build-arg CACHEBUST=`git rev-parse hpu` -f Dockerfile-hpu -t dockerfile-hpu . 
-	docker run -rm -it   --runtime=habana  -e PT_HPU_LAZY_MODE=$(lazy) -e HABANA_VISIBLE_DEVICES=all   -e OMPI_MCA_btl_vader_single_copy_mechanism=none --shm-size 50G   --cap-add=sys_nice   --net=host   dockerfile-hpu:latest bash -c '. $$MILABENCH_VENV/bin/activate && milabench install --use-current-env --select $(bench) && pip uninstall torch torchvision torchaudio -y  &&  milabench prepare --use-current-env --select $(bench) && milabench run --use-current-env $(args) --select $(bench)'
+	docker run --rm -it   --runtime=habana  -e PT_HPU_LAZY_MODE=$(lazy) -e HABANA_VISIBLE_DEVICES=all   -e OMPI_MCA_btl_vader_single_copy_mechanism=none --shm-size 50G   --cap-add=sys_nice   --net=host   dockerfile-hpu:latest bash -c '. $$MILABENCH_VENV/bin/activate && milabench install --use-current-env --select $(bench) && pip uninstall torch torchvision torchaudio -y  &&  milabench prepare --use-current-env --select $(bench) && milabench run --use-current-env $(args) --select $(bench)'