Skip to content

darts_utils.cuda

Utility functions around cuda, e.g. memory management.

Functions:

free_cupy 

free_cupy()

Free the CUDA memory of cupy.

Source code in darts-utils/src/darts_utils/cuda.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
def free_cupy():
    """Free the CUDA memory of cupy."""
    try:
        import cupy as cp
    except ImportError:
        cp = None

    if cp is not None:
        gc.collect()
        cp.get_default_memory_pool().free_all_blocks()
        cp.get_default_pinned_memory_pool().free_all_blocks()

free_torch 

free_torch()

Free the CUDA memory of pytorch.

Source code in darts-utils/src/darts_utils/cuda.py 
19
20
21
22
23
24
25
def free_torch():
    """Free the CUDA memory of pytorch."""
    import torch

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()