diff --git a/.github/workflows/gaudi2.yml b/.github/workflows/gaudi2.yml
index 746beab3ce42..cd4b2c6e55fe 100644
--- a/.github/workflows/gaudi2.yml
+++ b/.github/workflows/gaudi2.yml
@@ -18,15 +18,20 @@ jobs:
   unit-tests:
     # The type of runner that the job will run on
     runs-on: [self-hosted, intel, gaudi2]
+    container:
+      image: vault.habana.ai/gaudi-docker/1.14.0/ubuntu22.04/habanalabs/pytorch-installer-2.1.1:latest
+      ports:
+        - 80
+      options: --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host
 
     # Steps represent a sequence of tasks that will be executed as part of the job
     steps:
       # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
       - uses: actions/checkout@v3
 
-      - id: setup-venv
-        uses: ./.github/workflows/setup-venv
-
-      - name: Python environment
+      - name: Check container state
         run: |
-          pip list
+          ldd --version
+          hl-smi
+          python -c "import torch; print('torch:', torch.__version__, torch)"
+          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"