get_cuda_memory_info.py

import torch
import dataclasses
def cuda_memory_info(device = torch.device("cuda")):

    @dataclasses.dataclass
    class CudaMemoryInfo:
        device: torch.device
        total: int
        reserved: int
        allocated: int
        reserved_percent: int = dataclasses.field(init=False)
        allocated_percent: int = dataclasses.field(init=False)

        def __post_init__(self):
            self.reserved_percent = int(self.reserved * 100.0 // self.total)
            self.allocated_percent = int(self.allocated * 100.0 // self.total)

    return CudaMemoryInfo(
        device=device,
        total=torch.cuda.get_device_properties(device=device).total_memory,
        reserved=torch.cuda.memory_reserved(device=device),
        allocated=torch.cuda.memory_allocated(device=device),
    )