(*args, **kwargs)
| 66 | |
| 67 | @wraps(func) |
| 68 | def wrapped(*args, **kwargs): |
| 69 | with _ignore_torch_cuda_oom(): |
| 70 | return func(*args, **kwargs) |
| 71 | |
| 72 | # Clear cache and retry |
| 73 | torch.cuda.empty_cache() |
| 74 | with _ignore_torch_cuda_oom(): |
| 75 | return func(*args, **kwargs) |
| 76 | |
| 77 | # Try on CPU. This slows down the code significantly, therefore print a notice. |
| 78 | logger = logging.getLogger(__name__) |
| 79 | logger.info("Attempting to copy inputs of {} to CPU due to CUDA OOM".format(str(func))) |
| 80 | new_args = (maybe_to_cpu(x) for x in args) |
| 81 | new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()} |
| 82 | return func(*new_args, **new_kwargs) |
| 83 | |
| 84 | return wrapped |
nothing calls this directly
no test coverage detected