diff --git a/main.sh b/main.sh
index 032abdb..69737ac 100755
--- a/main.sh
+++ b/main.sh
@@ -5,7 +5,7 @@ chmod +x ./nvidia-pikaos-kernel-module/get_pwd.sh
 
 echo "$(apt-cache show nvidia-driver-$DRIVER | grep Version: | head -n1 | cut -f2 -d":" | cut -f1,2,3 -d"." | cut -f1 -d"-" | tr -d ' ')" > ./nvidia-pikaos-kernel-module/DRIVER
 echo "$(apt-cache show kernel-pika | grep Depends: | head -n1 | cut -f2 -d":" | cut -f1 -d"," | cut -f3,4 -d"-" | tr -d ' ')" > ./nvidia-pikaos-kernel-module/KERNEL
-echo "$(apt-cache show nvidia-dkms-535 | grep Version: | head -n1 | cut -f2 -d":" | tr -d ' ')" > ./nvidia-pikaos-kernel-module/DRIVER_VERSION
+echo "$(apt-cache show nvidia-kernel-source-535 | grep Version: | head -n1 | cut -f2 -d":" | tr -d ' ')" > ./nvidia-pikaos-kernel-module/DRIVER_VERSION
 
 cd ./nvidia-pikaos-kernel-module
 
@@ -21,10 +21,8 @@ DEBIAN_FRONTEND=noninteractive
 
 # Get build deps
 apt-get build-dep ./ -y
-#apt download nvidia-dkms-$DRIVER -y
-#dpkg-deb -x ./nvidia-dkms-$DRIVER*.deb /
-wget http://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/nvidia-kernel-dkms_535.104.05-1_amd64.deb
-dpkg-deb -x ./nvidia-kernel-dkms_535.104.05-1_amd64.deb /
+apt download nvidia-kernel-source-$DRIVER -y
+dpkg-deb -x ./nvidia-kernel-source-$DRIVER*.deb /
 
 # Build package
 dpkg-buildpackage --no-sign