Torch functional tensors needs sync'ing before transferring to CPU.

This patch fixes #8482.
pytorch · Dec 11, 2024 · 7a162ed · 7a162ed
1 parent 7dd2697
commit 7a162ed
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torch_xla/core/xla_model.py b/torch_xla/core/xla_model.py
@@ -1281,6 +1281,9 @@ def convert_fn(tensors):
         tensors, devices=[], wait=True, sync_xla_data=True)
     if not convert:
       return tensors
+    for t in tensors:
+      if torch._is_functional_tensor(t):
+        torch._functionalize_sync(t)
     return torch_xla._XLAC._xla_get_cpu_tensors(tensors)
 
   def select_fn(v):