add onnxruntime-traning

huggingface · Jan 10, 2025 · d001b9b · d001b9b
1 parent d0f62b0
commit d001b9b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/.github/workflows/test_onnxruntime_training.yml b/.github/workflows/test_onnxruntime_training.yml
@@ -51,7 +51,7 @@ jobs:
           pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
           pip install --no-cache-dir torch-ort onnxruntime-training && python -m torch_ort.configure
           pip install --no-cache-dir evaluate absl-py rouge_score seqeval sacrebleu nltk scikit-learn
-          pip install .[tests,onnxruntime-gpu]
+          pip install .[tests,onnxruntime-training]
 
       - name: Test with pytest
         run: |

diff --git a/setup.py b/setup.py
@@ -89,6 +89,10 @@
         "executorch>=0.4.0",
         "transformers>=4.46",
     ],
+    "onnxruntime-training": [
+        "torch-ort",
+        "onnxruntime-training",
+    ],
     "diffusers": ["diffusers<0.32.0"],
     "intel": "optimum-intel>=1.18.0",
     "openvino": "optimum-intel[openvino]>=1.18.0",