()
| 30 | |
| 31 | |
| 32 | def clean_memory(): |
| 33 | gc.collect() |
| 34 | if HAS_CUDA: |
| 35 | torch.cuda.empty_cache() |
| 36 | if HAS_XPU: |
| 37 | torch.xpu.empty_cache() |
| 38 | if HAS_MPS: |
| 39 | torch.mps.empty_cache() |
| 40 | |
| 41 | |
| 42 | def clean_memory_on_device(device: Optional[Union[str, torch.device]]): |