From 2ef583eb02f514ed21e9bd65b6d027086b7418ff Mon Sep 17 00:00:00 2001
From: Manfei <41607353+ManfeiBai@users.noreply.github.com>
Date: Thu, 31 Aug 2023 15:55:17 -0700
Subject: [PATCH] Update llama2-model.libsonnet

---
 tests/pytorch/nightly/llama2-model.libsonnet | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/pytorch/nightly/llama2-model.libsonnet b/tests/pytorch/nightly/llama2-model.libsonnet
index 8b37a251a..3f34832c7 100644
--- a/tests/pytorch/nightly/llama2-model.libsonnet
+++ b/tests/pytorch/nightly/llama2-model.libsonnet
@@ -152,6 +152,8 @@ local utils = import 'templates/utils.libsonnet';
         export PROFILE_EPOCH=2
         export PROFILE_STEP=0
         export PROFILE_DURATION_MS=20000
+        export XLA_USE_SPMD=1
+        export PJRT_DEVICE=TPU
       |||,
       tpuVmExtraSetup: |||
         pip3 uninstall torch torch_xla torchvision libtpu-nightly -y