diff --git a/cog.yaml b/cog.yaml
index d835b67..82d01c7 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -6,7 +6,6 @@ build:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
   python_packages:
-    - "torch==1.10.1"
-    - "flax==0.5.2"
+    - "min-dalle>=0.2.10"
 
 predict: "replicate/predict.py:Predictor"