diff --git a/benchmarks/lightning/main.py b/benchmarks/lightning/main.py
index 165a0ad2..a9bdcba0 100644
--- a/benchmarks/lightning/main.py
+++ b/benchmarks/lightning/main.py
@@ -87,17 +87,14 @@ def main():
     observer, monitor = prepare_voir()
     loader = observer.loader(imagenet_dataloader(args, model, rank, world_size))
 
-    # alueError: The `HPUAccelerator` can only be used with a `SingleHPUStrategy` or `HPUParallelStrategy`, found DDPStrategy.
-    
-    # from pytorch_lightning.strategies import HPUParallelStrategy
-    from lightning_habana import HPUParallelStrategy
+    # ValueError: The `HPUAccelerator` can only be used with a `SingleHPUStrategy` or `HPUParallelStrategy`, found DDPStrategy.
     
     # train model
     trainer = L.Trainer(
-        accelerator="hpu", 
+        accelerator="auto", 
         devices=n, 
         num_nodes=nnodes, 
-        strategy=HPUParallelStrategy(),
+        strategy="auto",
         max_epochs=args.epochs,
         precision="bf16-mixed",
         enable_checkpointing=False,
diff --git a/docker/Makefile b/docker/Makefile
index d149f4e4..48333830 100644
--- a/docker/Makefile
+++ b/docker/Makefile
@@ -1,7 +1,8 @@
 
 
 
-bench = rlhf-single
+# bench = rlhf-single
+bench = lightning
 
 hpu:
 	git add --all