set PJRT_DEVICE before loading torch_xla

Mitigation for #326 PiperOrigin-RevId: 692234642
google-ai-edge · Nov 1, 2024 · bd09407 · bd09407
1 parent 284fe62
commit bd09407
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/ai_edge_torch/lowertools/torch_xla_utils.py b/ai_edge_torch/lowertools/torch_xla_utils.py
@@ -19,9 +19,14 @@
 import gc
 import itertools
 import logging
+import os
 import tempfile
 from typing import Any, Dict, Optional, Tuple, Union
 
+if "PJRT_DEVICE" not in os.environ:
+  # https://github.com/google-ai-edge/ai-edge-torch/issues/326
+  os.environ["PJRT_DEVICE"] = "CPU"
+
 from ai_edge_torch import model
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch._convert import signature as signature_module