diff --git a/gantry-train-requirements.txt b/gantry-train-requirements.txt
index 441796f..e65d726 100644
--- a/gantry-train-requirements.txt
+++ b/gantry-train-requirements.txt
@@ -1,3 +1,4 @@
+torchvision
 accelerate
 peft
 wandb