Skip to content

Commit 15cf7eb

Browse files
committed
Added a potential solution for windows
1 parent d99f183 commit 15cf7eb

File tree

2 files changed

+28
-1
lines changed

2 files changed

+28
-1
lines changed

py/torch_tensorrt/dynamo/conversion/_conversion.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,6 @@ def convert_module(
121121
with io.BytesIO() as engine_bytes:
122122
engine_bytes.write(serialized_engine)
123123
serialized_engine = engine_bytes.getvalue()
124-
breakpoint()
125124
return rt_cls(
126125
serialized_engine=serialized_engine,
127126
input_binding_names=list(interpreter_result.input_names),

py/torch_tensorrt/dynamo/utils.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -881,3 +881,31 @@ def release_memory() -> None:
881881
logger.warning("Failed to release CPU memory.")
882882
except Exception:
883883
logger.warning("Failed to release CPU memory.")
884+
885+
elif platform.system() == "Windows":
886+
try:
887+
from ctypes import wintypes as wt
888+
889+
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
890+
psapi = ctypes.WinDLL("psapi", use_last_error=True)
891+
892+
GetCurrentProcess = kernel32.GetCurrentProcess
893+
GetCurrentProcess.restype = wt.HANDLE
894+
hProcess = GetCurrentProcess()
895+
896+
EmptyWorkingSet = psapi.EmptyWorkingSet
897+
EmptyWorkingSet.argtypes = [wt.HANDLE]
898+
EmptyWorkingSet.restype = wt.BOOL
899+
EmptyWorkingSet(hProcess)
900+
901+
SetProcessWorkingSetSize = kernel32.SetProcessWorkingSetSize
902+
SetProcessWorkingSetSize.argtypes = [
903+
wt.HANDLE,
904+
ctypes.c_size_t,
905+
ctypes.c_size_t,
906+
]
907+
SetProcessWorkingSetSize.restype = wt.BOOL
908+
SIZE_T_MAX = ctypes.c_size_t(-1).value # 0xFFFFFFFFFFFFFFFF on x64
909+
SetProcessWorkingSetSize(hProcess, SIZE_T_MAX, SIZE_T_MAX)
910+
except Exception:
911+
logger.warning("Failed to release CPU memory.")

0 commit comments

Comments
 (0)