[sleep mode] clear pytorch cache after sleep (#15248)
Signed-off-by: <villard@us.ibm.com>
This commit is contained in:
@@ -8,6 +8,7 @@
|
||||
# not sure why, they are created from a different context.
|
||||
# the only successful approach is to call cuda driver API in C.
|
||||
import dataclasses
|
||||
import gc
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Callable, Dict, Optional, Tuple, Union
|
||||
@@ -204,6 +205,9 @@ class CuMemAllocator:
|
||||
data.cpu_backup_tensor = cpu_backup_tensor
|
||||
unmap_and_release(handle)
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def wake_up(self):
|
||||
"""
|
||||
Wake up the allocator from sleep mode.
|
||||
|
||||
Reference in New Issue
Block a user