nemo_curator.stages.resources

Module Contents

Name	Description
`Resources`	Define resource requirements for a processing stage.

Name	Description
`_get_gpu_memory_gb`	Get GPU memory in GB for the current device.

class nemo_curator.stages.resources.Resources(
    cpus: float = 1.0,
    gpu_memory_gb: float = 0.0,
    gpus: float = 0.0
)

Dataclass

Define resource requirements for a processing stage.

cpus

float = 1.0

gpu_memory_gb

float = 0.0

gpus

float = 0.0

requires_gpu

bool

Check if this stage requires GPU resources.

nemo_curator.stages.resources.Resources.__post_init__()

Calculate GPU count based on memory requirements.

nemo_curator.stages.resources._get_gpu_memory_gb() -> float

Get GPU memory in GB for the current device.