# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/gpu/torch.ipynb.

# %% auto 0
__all__ = ['get_device', 'cuda_gpu_processes', 'clean_gpu_cache']

# %% ../../nbs/gpu/torch.ipynb 3
from typing import Optional
import gc

import torch

# %% ../../nbs/gpu/torch.ipynb 5
from ..seed import set_torch as set_torch_seed

# %% ../../nbs/gpu/torch.ipynb 7
def get_device(selected_gpu_idx: Optional[int]):
    if torch.cuda.is_available():
        if selected_gpu_idx is None:
            selected_gpu_idx = torch.cuda.current_device()
        return torch.device(f'cuda:{selected_gpu_idx}')
    if torch.backends.mps.is_available():
        return torch.device('mps')
    return torch.device('cpu')

# %% ../../nbs/gpu/torch.ipynb 10
def cuda_gpu_processes(device: Optional[torch.device] = None):
    if device is None:
        device = get_device()
    if device.type == 'cuda':
        return torch.cuda.list_gpu_processes(device)
    return None


# %% ../../nbs/gpu/torch.ipynb 13
def clean_gpu_cache() -> bool:
    """Clean GPU cache and return True if GPU is available."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        return True
    return False
