[PT/XLA] Add missing deps for CI test (#189)

GoogleCloudPlatform · Mar 8, 2024 · 421e856 · 421e856
1 parent 264cd87
commit 421e856
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/dags/legacy_test/tests/pytorch/nightly/ci.libsonnet b/dags/legacy_test/tests/pytorch/nightly/ci.libsonnet
@@ -17,14 +17,16 @@ local tpus = import 'templates/tpus.libsonnet';
         test/tpu/run_tests.sh
       |||,
     ],
+
+  },
+  local pjrt = self.pjrt,
+  pjrt:: common.PyTorchTpuVmMixin {
     tpuSettings+: {
       tpuVmExtraSetup: |||
-        pip install expecttest==0.1.6
+        pip install expecttest==0.1.6 rich
       |||,
     },
   },
-  local pjrt = self.pjrt,
-  pjrt:: common.PyTorchTpuVmMixin,
 
   local v5litepod_4 = self.v5litepod_4,
   v5litepod_4:: {