Spaces:
Running
Running
File size: 472 Bytes
e0336bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
import torch
def clean_memory_on_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: torch.device):
if device.type == "cuda":
torch.cuda.synchronize()
elif device.type == "xpu":
torch.xpu.synchronize()
elif device.type == "mps":
torch.mps.synchronize()
|