diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh
index b480cee78..5b7a65e80 100755
--- a/fastvideo-kernel/build.sh
+++ b/fastvideo-kernel/build.sh
@@ -11,11 +11,6 @@ echo "Building fastvideo-kernel..."
 # Install build dependencies
 pip install scikit-build-core cmake ninja
 
-# Set TORCH_CUDA_ARCH_LIST only if not already set
-if [ -z "$TORCH_CUDA_ARCH_LIST" ]; then
-    export TORCH_CUDA_ARCH_LIST="9.0a"
-fi
-
 # Build and install
 # Use -v for verbose output
 pip install . -v --no-build-isolation
diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml
index fc043c8f2..8eb61c00b 100644
--- a/fastvideo-kernel/pyproject.toml
+++ b/fastvideo-kernel/pyproject.toml
@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.0"
+version = "0.2.1"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -32,4 +32,4 @@ dependencies = [
 [tool.scikit-build]
 cmake.build-type = "Release"
 minimum-version = "build-system.requires"
-wheel.packages = ["python"]
+wheel.packages = ["python/fastvideo_kernel"]
diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py
index cd4a74f1d..22ede949e 100644
--- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py
+++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.1.0"
+from .version import __version__
 
 from fastvideo_kernel.ops import (
     sliding_tile_attention,
diff --git a/fastvideo-kernel/python/fastvideo_kernel/version.py b/fastvideo-kernel/python/fastvideo_kernel/version.py
new file mode 100644
index 000000000..3ced3581b
--- /dev/null
+++ b/fastvideo-kernel/python/fastvideo_kernel/version.py
@@ -0,0 +1 @@
+__version__ = "0.2.1"