diff --git a/.github/workflows/pypi-publish.yaml b/.github/workflows/pypi-publish.yaml
index 14a1e73..602c7d4 100644
--- a/.github/workflows/pypi-publish.yaml
+++ b/.github/workflows/pypi-publish.yaml
@@ -26,6 +26,11 @@ jobs:
           rm -rf dist
           python -m build
 
+      - name: Quick test of the built dist
+        run: |
+          pip install "src/python/dist/$(ls src/python/dist/ | head -1)"
+          python -c "import pose_format"
+
       - name: Publish distribution 📦 to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/python/__init__.py b/src/python/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/python/pyproject.toml b/src/python/pyproject.toml
index d3601a0..1d5706c 100644
--- a/src/python/pyproject.toml
+++ b/src/python/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "pose_format"
 description = "Library for viewing, augmenting, and handling .pose files"
-version = "0.2.0"
+version = "0.2.1"
 keywords = ["Pose Files", "Pose Interpolation", "Pose Augmentation"]
 authors = [
     { name = "Amit Moryossef", email = "amitmoryossef@gmail.com" }
@@ -26,9 +26,20 @@ dev = [
 
 [tool.setuptools]
 packages = [
-    "pose_format"
+    "pose_format",
+    "pose_format.utils",
+    "pose_format.testing",
+    "pose_format.numpy",
+    "pose_format.numpy.representation",
+    "pose_format.torch",
+    "pose_format.torch.masked",
+    "pose_format.torch.representation",
+    "pose_format.tensorflow",
+    "pose_format.tensorflow.masked",
+    "pose_format.tensorflow.representation"
 ]
 
+
 [tool.pytest.ini_options]
 addopts = "-v"
 filterwarnings = [
@@ -36,6 +47,7 @@ filterwarnings = [
 ]
 testpaths = [
     "pose_format",
+    "tests"
 ]
 
 [tool.yapf]
diff --git a/src/python/pose_format/pose_test.py b/src/python/tests/pose_test.py
similarity index 100%
rename from src/python/pose_format/pose_test.py
rename to src/python/tests/pose_test.py
diff --git a/src/python/pose_format/pose_tf_graph_mode_test.py b/src/python/tests/pose_tf_graph_mode_test.py
similarity index 100%
rename from src/python/pose_format/pose_tf_graph_mode_test.py
rename to src/python/tests/pose_tf_graph_mode_test.py