clean_memory.py 211 B

12345678910
  1. # Copyright (c) Opendatalab. All rights reserved.
  2. import torch
  3. import gc
  4. def clean_memory():
  5. if torch.cuda.is_available():
  6. torch.cuda.empty_cache()
  7. torch.cuda.ipc_collect()
  8. gc.collect()