[diffusion] Clean code (#19325)

2026-02-25 21:16:03 +03:00
parent 2fb239450e
commit 0217e82a08
8 changed files with 55 additions and 40 deletions
@@ -90,9 +90,9 @@ def _torch_cleanup() -> None:
    try:
        import torch

-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-            torch.cuda.empty_cache()
+        if torch.get_device_module().is_available():
+            torch.get_device_module().synchronize()
+            torch.get_device_module().empty_cache()
    except Exception:
        pass