File tree Expand file tree Collapse file tree 2 files changed +2
-31
lines changed Expand file tree Collapse file tree 2 files changed +2
-31
lines changed Original file line number Diff line number Diff line change @@ -70,7 +70,7 @@ class UnsupportedOperatorException(RuntimeError):
70
70
71
71
72
72
class TRTInterpreterResult (NamedTuple ):
73
- engine : trt .ICudaEngine
73
+ engine : trt .ICudaEngine | bytes
74
74
input_names : Sequence [str ]
75
75
output_names : Sequence [str ]
76
76
weight_name_map : Optional [dict [Any , Any ]]
@@ -596,7 +596,7 @@ def _save_weight_mapping(self) -> None:
596
596
torch .cuda .empty_cache ()
597
597
598
598
@needs_refit # type: ignore[misc]
599
- def _insert_engine_to_cache (self , hash_val : str , engine : bytes ) -> None :
599
+ def _insert_engine_to_cache (self , hash_val : str , engine : trt . ICudaEngine ) -> None :
600
600
serialized_engine = engine .serialize ()
601
601
# TODO: @Evan is waiting for TRT's feature to cache the weight-stripped engine
602
602
# if not self.compilation_settings.strip_engine_weights:
Original file line number Diff line number Diff line change @@ -882,32 +882,3 @@ def release_memory() -> None:
882
882
logger .warning ("Failed to release CPU memory." )
883
883
except Exception :
884
884
logger .warning ("Failed to release CPU memory." )
885
-
886
- elif platform .system () == "Windows" :
887
- from ctypes import wintypes
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 = wintypes .HANDLE
894
- hproc = GetCurrentProcess ()
895
-
896
- HeapSetInformation = kernel32 .HeapSetInformation
897
- HeapSetInformation .argtypes = [
898
- wintypes .HANDLE ,
899
- ctypes .c_int ,
900
- ctypes .c_void_p ,
901
- ctypes .c_size_t ,
902
- ]
903
- HeapSetInformation .restype = wintypes .BOOL
904
- GetProcessHeap = kernel32 .GetProcessHeap
905
- GetProcessHeap .restype = wintypes .HANDLE
906
- ok = False
907
- try :
908
- HeapOptimizeResources = 3
909
- hheap = GetProcessHeap ()
910
- if HeapSetInformation (hheap , HeapOptimizeResources , None , 0 ):
911
- ok = True
912
- except Exception :
913
- logger .warning ("Failed to release CPU memory." )
You can’t perform that action at this time.
0 commit comments