Enable megacore_dense by default (#5520)

Summary: This change enables megacore_dense by default to allow asynchorous cc ops especailly for GSPMD. Test Plan: CI
pytorch · Aug 30, 2023 · 17a0b7c · 17a0b7c
1 parent 17513f7
commit 17a0b7c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -40,6 +40,7 @@ def _setup_default_env():
   _set_missing_env('GRPC_VERBOSITY', 'ERROR')
   _set_missing_env('ALLOW_MULTIPLE_LIBTPU_LOAD', '1')
   _set_missing_env('TPU_ML_PLATFORM', 'PyTorch/XLA')
+  _set_missing_env('TPU_MEGACORE', 'megacore_dense')
 
 
 _fd, _tmp_fname = -1, ''