From c269e70541ff89aa73b5521956bb4e3f8b83902c Mon Sep 17 00:00:00 2001
From: fangyuan-3149 <314945674@qq.com>
Date: Fri, 28 Aug 2026 15:07:17 +0800
Subject: [PATCH] fix(engine): handle torch 2.9 allocator API deprecation

---
 python/freetoken/engine/engine.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index cd6505d2..bccccf79 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -1024,7 +1024,11 @@ def _ensure_expandable_segments() -> None:
     if os.environ.get("PYTORCH_ALLOC_CONF") or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"):
         return
     try:
-        torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+        # torch 2.9+: _set_allocator_settings -> _C._cuda_setAllocatorSettings
+        try:
+            torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+        except AttributeError:
+            torch._C._cuda_setAllocatorSettings("expandable_segments:True")
     except Exception as exc:  # pragma: no cover - depends on torch build
         logger.info_rank0(f"Could not enable expandable_segments ({exc}); continuing")
         return