diff --git a/.github/workflows/docker-build-and-run.yml b/.github/workflows/docker-build-and-run.yml
new file mode 100644
index 0000000..e925c74
--- /dev/null
+++ b/.github/workflows/docker-build-and-run.yml
@@ -0,0 +1,78 @@
+name: Build and Run Docker
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [opened, synchronize, edited]
+    branches:
+      - main
+      - "docker/**"
+jobs:
+  build-and-run-docker:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.x'
+
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+
+      - name: Generate input files
+        run: |
+          python -m Docker.generate_signal_docker_test
+
+      - name: Verify input files
+        run: |
+          for file in ivim_simulation.nii.gz ivim_simulation.bval ivim_simulation.bvec; do
+            if [ ! -f "$file" ]; then
+              echo "Error: $file not found"
+              exit 1
+            fi
+          done
+          echo "All input files generated successfully"
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build Docker image
+        run: |
+          docker build -t tf2.4_ivim-mri_codecollection -f Docker/Dockerfile .
+
+      - name: Run Docker container
+        run: |
+          docker run --rm --name TF2.4_IVIM-MRI_CodeCollection \
+            -v ${{ github.workspace }}:/usr/src/app \
+            -v ${{ github.workspace }}:/usr/app/output \
+            tf2.4_ivim-mri_codecollection ivim_simulation.nii.gz ivim_simulation.bvec ivim_simulation.bval
+
+      - name: Verify output files
+        run: |
+          for file in f.nii.gz dp.nii.gz d.nii.gz; do
+            if [ ! -f "$file" ]; then
+              echo "Error: $file not found"
+              exit 1
+            fi
+          done
+          echo "All output files generated successfully"
+
+      - name: Clean up artifacts and Docker image
+        run: |
+          docker rmi tf2.4_ivim-mri_codecollection || true
+          rm -f tf2.4_ivim-mri_codecollection.tar.gz
+          rm -f ${{ github.workspace }}/f.nii.gz
+          rm -f ${{ github.workspace }}/dp.nii.gz
+          rm -f ${{ github.workspace }}/d.nii.gz
+          rm -f ${{ github.workspace }}/ivim_simulation.nii.gz
+          rm -f ${{ github.workspace }}/ivim_simulation.bval
+          rm -f ${{ github.workspace }}/ivim_simulation.bvec
+      - name: Cleanup Docker
+        run: |
+          docker system prune -a --force
diff --git a/Docker/generate_signal_docker_test.py b/Docker/generate_signal_docker_test.py
new file mode 100644
index 0000000..c6c96a3
--- /dev/null
+++ b/Docker/generate_signal_docker_test.py
@@ -0,0 +1,42 @@
+import numpy as np
+import nibabel as nib
+from utilities.data_simulation.GenerateData import GenerateData
+from WrapImage.nifti_wrapper import save_nifti_file
+
+
+def save_bval_bvec(filename, values):
+    if filename.endswith('.bval'):
+        # Convert list to a space-separated string for bval
+        values_string = ' '.join(map(str, values))
+    elif filename.endswith('.bvec'):
+        # Convert 2D list to a line-separated, space-separated string for bvec
+        values_string = '\n'.join(' '.join(map(str, row)) for row in values)
+    else:
+        raise ValueError("Unsupported file extension. Use '.bval' or '.bvec'.")
+
+    with open(filename, 'w') as file:
+        file.write(values_string)
+
+# Set random seed for reproducibility
+np.random.seed(42)
+# Create GenerateData instance
+gd = GenerateData()
+
+# Generate random input data
+shape = (10, 10, 5)
+f_in = np.random.uniform(low=0, high=1, size=shape)
+D_in = np.random.uniform(low=0, high=1e-3, size=shape)
+Dp_in = np.random.uniform(low=0, high=1e-1, size=shape)
+S0 = 1000  # Setting a constant S0 for simplicity
+bvals = np.array([0, 50, 100, 500, 1000])
+bvals_reshaped = np.broadcast_to(bvals, shape)
+
+# Generate IVIM signal
+signals = gd.ivim_signal(D_in, Dp_in, f_in, S0, bvals_reshaped)
+
+# Save the generated image as a NIfTI file
+save_nifti_file(signals, "ivim_simulation.nii.gz")
+# Save the bval in a file
+save_bval_bvec("ivim_simulation.bval", [0, 50, 100, 500, 1000])
+# Save the bvec value 
+save_bval_bvec("ivim_simulation.bvec", [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
diff --git a/WrapImage/nifti_wrapper.py b/WrapImage/nifti_wrapper.py
index 5eae67e..6f43792 100644
--- a/WrapImage/nifti_wrapper.py
+++ b/WrapImage/nifti_wrapper.py
@@ -56,7 +56,9 @@ def save_nifti_file(data, output_file, affine=None, **kwargs):
     For saving the 3d nifti images of the output of the algorithm
     """
     if affine is None:
-        affine = np.eye(data.ndim + 1)
+        affine = np.eye(4)
+    else:
+        affine = np.array(affine.reshape(4, 4))
     output_img = nib.nifti1.Nifti1Image(data, affine , **kwargs)
     nib.save(output_img, output_file)
 
diff --git a/requirements.txt b/requirements.txt
index 8964af8..29742e3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 numpy<2
+nibabel
 scipy
 torchio
 torch
diff --git a/utilities/data_simulation/GenerateData.py b/utilities/data_simulation/GenerateData.py
index f8980a4..f1309e2 100644
--- a/utilities/data_simulation/GenerateData.py
+++ b/utilities/data_simulation/GenerateData.py
@@ -53,7 +53,7 @@ def exponential_signal(self, D, bvalues):
         bvalues : list or array of float
             The diffusion (b-values)
         """
-        assert D >= 0, 'D must be >= 0'
+        assert np.all(D >= 0), 'all values in D must be >= 0'
         return self._op.exp(-self._op.asarray(bvalues, dtype='float64') * D)
 
     def multiexponential_signal(self, D, F, S0, bvalues):
@@ -116,7 +116,7 @@ def linear_signal(self, D, bvalues, offset=0):
         offset : float
             The signal offset
         """
-        assert D >= 0, 'D must be >= 0'
+        assert np.all(D >= 0), 'every value in D must be >= 0'
         data = -D * np.asarray(bvalues)
         return data + offset
 
@@ -144,4 +144,4 @@ def multilinear_signal(self, D, F, S0, bvalues, offset=0):
             signal += f * self.linear_signal(d, bvalues)
         signal *= S0
         signal += offset
-        return signal
\ No newline at end of file
+        return signal