Set XLA_USE_SPMD for spmd cpp tests. (#6273)

pytorch · Jan 11, 2024 · a728afe · a728afe
1 parent 68f4750
commit a728afe
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/test/cpp/test_xla_sharding.cpp b/test/cpp/test_xla_sharding.cpp
@@ -33,7 +33,13 @@ bool XlaDataValuesEqual(torch::lazy::BackendDataPtr a,
 }
 }  // namespace
 
-class XLAShardingTest : public AtenXlaTensorTestBase {};
+class XLAShardingTest : public AtenXlaTensorTestBase {
+ protected:
+  static void SetUpTestCase() {
+    setenv("XLA_USE_SPMD", "1", /*overwrite=*/true);
+    CommonSetup();
+  }
+};
 
 TEST_F(XLAShardingTest, GetShardShape) {
   auto tensor = at::ones({8, 7}, at::TensorOptions(at::kFloat));