`nemo_rl.utils.memory_tracker`#

Module Contents#

Classes#

`MemoryTrackerDataPoint`
`MemoryTracker`

API#

class nemo_rl.utils.memory_tracker.MemoryTrackerDataPoint#

Bases: pydantic.BaseModel

stage: str#: None

memory_used_before_stage_gb: float#: None

variables_before_stage: List[str]#: None

memory_used_after_stage_gb: Optional[float]#: None

variables_after_stage: Optional[List[str]]#: None

property mem_used_diff_gb: float#

property new_variables: List[str]#

get_snapshot_str() → str#

class nemo_rl.utils.memory_tracker.MemoryTracker#

Bases: pydantic.BaseModel

data_points: List[nemo_rl.utils.memory_tracker.MemoryTrackerDataPoint]#: ‘Field(…)’

model_post_init(context)#

snapshot_start_of_stage( new_stage: str, all_current_variables: List[str], ) → None#