clear pending ir should also clear the cc op tokens (#5385)

pytorch · Aug 1, 2023 · 2676a20 · 2676a20
1 parent c0a432b
commit 2676a20
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1729,6 +1729,8 @@ void InitXlaModuleBindings(py::module m) {
     // Use with caution. Those tensor whole ir was cleared with be replaced
     // with a placeholder XLAData and SHOULD NOT be accessed.
     ClearPendingIrs(device);
+    auto xla_device = GetDeviceOrCurrent(device);
+    SetAllReduceToken(xla_device, nullptr);
   });
 
   m.def("_run_cached_graph",