mirror of
https://git.datalinker.icu/ali-vilab/TeaCache
synced 2025-12-08 20:34:24 +08:00
13 lines
207 B
Python
13 lines
207 B
Python
import functools
|
|
|
|
import torch
|
|
|
|
|
|
def empty_cache(func):
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
torch.cuda.empty_cache()
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|