From 1dafbf1247c1984e1af0fc4e7b90364ffbfe6fe3 Mon Sep 17 00:00:00 2001
From: Michael Benayoun <mickbenayoun@gmail.com>
Date: Thu, 3 Aug 2023 13:51:17 +0200
Subject: [PATCH] Add neuronx

---
 .github/workflows/test_on_trainium.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/test_on_trainium.yml b/.github/workflows/test_on_trainium.yml
index 87a385d16..848a3506f 100644
--- a/.github/workflows/test_on_trainium.yml
+++ b/.github/workflows/test_on_trainium.yml
@@ -63,7 +63,7 @@ jobs:
       # - name: Install python3.8-venv
       #   run: sudo apt update; sudo apt install -y python3.8-venv
       - name: Install Python dependencies
-        run: pip install .[tests]
+        run: pip install .[tests,neuronx]
       - name: Run tests on neuron cores
         run: |
           HF_TOKEN_OPTIMUM_NEURON_CI=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} USE_VENV="false" pytest -m "is_trainium_test" tests -k "not test_run_example"