Working on unit tests workflow YAML, wip

Author: jbieniusiewi

.github/workflows/unit_test.yml

@@ -4,30 +4,33 @@ on:
   push:
     branches:
       - main
+      - unit_tests_in_ci
   pull_request:
     branches:
       - main
+      - unit_tests_in_ci
+  workflow_dispatch:
 
 jobs:
   unit_test_pytorch_2_4_1_cuda_12_1:
     runs-on: ubuntu-latest
 
     container:
-      image: pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime
+      image: pytorch/pytorch:2.4.1-cuda12.1-cudnn9-devel
 
     steps:
       - name: Checkout code
         uses: actions/checkout@v3
 
-      - name: Set up Python environment
+      - name: Set up environment
         run: |
           python --version
           pip install --upgrade pip
 
-      - name: Install dependencies
+      - name: Install package
         run: |
+          pip install pytest lightning
           pip install .
-          pip install pytest
 
       - name: Run Fault Tolerance unit tests
         run: |