diff --git a/Makefile b/Makefile
index ff8e5aa..c8ba2a8 100644
--- a/Makefile
+++ b/Makefile
@@ -7,8 +7,8 @@ all: $(targets)
 
 .PHONY: setup lock
 setup:
-	conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=12.1 -c pytorch -c nvidia
 	conda install -c nvidia cuda-toolkit=12.1
+	conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
 	pip install packaging
 	pip install flash-attn --no-build-isolation
 	pip install -U pip