From 017382841d6df6a8704bce2240fe25f193dee5e8 Mon Sep 17 00:00:00 2001
From: Dmitry Nikolaev <dmitry.nikolaev@amd.com>
Date: Wed, 24 Jun 2026 17:03:42 +0000
Subject: [PATCH] [ROCm] sync on pytorch module exit

---
 torch/csrc/cuda/Module.cpp | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/torch/csrc/cuda/Module.cpp b/torch/csrc/cuda/Module.cpp
index e53d5ea269b9e..6181ec02d7ed0 100644
--- a/torch/csrc/cuda/Module.cpp
+++ b/torch/csrc/cuda/Module.cpp
@@ -2285,6 +2285,16 @@ void initModule(PyObject* module) {
   registerCudaDeviceProperties(module);
   registerCudaPluggableAllocator(module);
   initCudaMethodBindings(module);
+#if defined(USE_ROCM)
+  // ROCm 7.x: drain the current device before C++ static destructors run.
+  // C++ std::atexit is too late; Python atexit matches torch.cuda.synchronize().
+  auto atexit_mod = py::module_::import("atexit");
+  atexit_mod.attr("register")(py::cpp_function([]() {
+    if (c10::cuda::device_count() > 0) {
+      c10::cuda::device_synchronize();
+    }
+  }));
+#endif
 }
 
 } // namespace torch::cuda