diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c6d904c..e2b768e 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -41,7 +41,9 @@ jobs:
 
     - name: Run tests with pytest and generate coverage
       run: |
-        pytest --cov=. --cov-report=xml --cov-report=term-missing
+        export PYTHONPATH=$PYTHONPATH:$(pwd)
+        echo $PYTHONPATH
+        python -m pytest --cov=. --cov-report=xml --cov-report=term-missing
 
     - name: Upload coverage report
       uses: actions/upload-artifact@v4
diff --git a/training/pyproject.toml b/training/pyproject.toml
index dceea6a..c7c62c0 100644
--- a/training/pyproject.toml
+++ b/training/pyproject.toml
@@ -5,3 +5,9 @@ line-length = 120
 [tool.isort]
 profile = "black"
 known_third_party = "wandb"
+
+[project]
+name = "alphafind_training"
+version = "0.0.1"
+description = "AlphaFold training -- setup for the similarity search on vast protein data"
+readme = "README.md"
diff --git a/training/tests/test_create_embedding.py b/training/tests/test_create_embedding.py
index cbfa8b3..66ecead 100644
--- a/training/tests/test_create_embedding.py
+++ b/training/tests/test_create_embedding.py
@@ -5,7 +5,7 @@
 
 import pandas as pd
 import pytest
-from alphafind_training.create_embedding import run
+from alphafind_training.create_embedding import create_embedding
 
 
 @pytest.fixture(scope="function")
@@ -37,7 +37,7 @@ def test_create_embedding():
         # 45 features for each protein - (10x10 - 10) / 2
         expected_dimensionality = 45
 
-        run(Path(cif_path), Path(output_path), granularity)
+        create_embedding(Path(cif_path), Path(output_path), granularity)
 
         assert os.path.exists(output_path)
         assert os.path.getsize(output_path) > 0