fixing python library requirements

kubeflow · Jan 4, 2024 · c6467de · c6467de
1 parent bf04ddc
commit c6467de
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
@@ -31,6 +31,11 @@
     "retrying>=1.3.3",
     "boto3>=1.33.9",
     "transformers>=4.35.2",
+    "torch>=2.1.1",
+    "torchvision>=0.16.1",
+    "torchaudio>=2.1.1",
+    "einops>=0.7.0",
+    "transformers_stream_generator>=0.0.4",
     "peft>=0.7.0",
 ]