diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
index ac3eb17..c97ead6 100644
--- a/.github/workflows/train_workflow.yml
+++ b/.github/workflows/train_workflow.yml
@@ -16,32 +16,50 @@ jobs:
         with:
           python-version: '3.10'
 
+      # Cache pip packages
+      - name: Cache pip packages
+        uses: actions/cache@v3
+        id: pip-cache
+        with:
+          path: |
+            ~/.cache/pip
+          key: ${{ runner.os }}-pip-v1
+          restore-keys: |
+            ${{ runner.os }}-pip-
+
+      # Cache PyTorch specific directories
       - name: Cache PyTorch
         uses: actions/cache@v3
         id: pytorch-cache
         with:
           path: |
-            ~/.cache/pip
             ~/.cache/torch
             ~/.cache/torch_extensions
-          key: ${{ runner.os }}-pytorch-${{ hashFiles('requirements.txt') }}-v1
+            ~/.local/lib/python*/site-packages/torch
+            ~/.local/lib/python*/site-packages/torch_*
+          key: ${{ runner.os }}-pytorch-v1
           restore-keys: |
             ${{ runner.os }}-pytorch-
 
+      # Install dependencies
       - name: Install dependencies
         run: |
-          pip install torch
-          pip install numpy
+          python -m pip install --upgrade pip
+          pip install torch numpy
+          # Print versions for debugging
+          python -c "import torch; print(f'PyTorch version: {torch.__version__}')"
+          python -c "import numpy; print(f'NumPy version: {numpy.__version__}')"
 
       - name: Create and run training script
         run: |
           echo "${{ inputs.script_content }}" > train.py
-          cat train.py  # Debug: print the content
+          echo "Content of train.py:"
+          cat train.py
           python train.py > training.log 2>&1
         
       - name: Upload logs
         uses: actions/upload-artifact@v3
-        if: always()  # Upload logs whether the job succeeds or fails
+        if: always()
         with:
           name: training-logs
-          path: training.log
\ No newline at end of file
+          path: training.log