A context which ignores CUDA OOM exception from pytorch.
()
| 10 | |
| 11 | @contextmanager |
| 12 | def _ignore_torch_cuda_oom(): |
| 13 | """ |
| 14 | A context which ignores CUDA OOM exception from pytorch. |
| 15 | """ |
| 16 | try: |
| 17 | yield |
| 18 | except RuntimeError as e: |
| 19 | # NOTE: the string may change? |
| 20 | if "CUDA out of memory. " in str(e): |
| 21 | pass |
| 22 | else: |
| 23 | raise |
| 24 | |
| 25 | |
| 26 | def retry_if_cuda_oom(func): |