29 lines
781 B
Python
29 lines
781 B
Python
from typing import Optional, Union
|
|
import torch
|
|
|
|
|
|
def clean_memory_on_device(device: Optional[Union[str, torch.device]]):
|
|
if device is None:
|
|
return
|
|
if isinstance(device, str):
|
|
device = torch.device(device)
|
|
if device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
elif device.type == "cpu":
|
|
pass
|
|
elif device.type == "mps": # not tested
|
|
torch.mps.empty_cache()
|
|
|
|
|
|
def synchronize_device(device: Optional[Union[str, torch.device]]):
|
|
if device is None:
|
|
return
|
|
if isinstance(device, str):
|
|
device = torch.device(device)
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize()
|
|
elif device.type == "xpu":
|
|
torch.xpu.synchronize()
|
|
elif device.type == "mps":
|
|
torch.mps.synchronize()
|