`nemo_rl.models.generation.megatron.utils`#

Module Contents#

`resolve_torch_dtype`	Convert a value to `torch.dtype`.
`log_gpu_memory`	Print a one-line GPU-memory summary for the calling rank.

nemo_rl.models.generation.megatron.utils.resolve_torch_dtype(val)#: Convert a value to torch.dtype.

nemo_rl.models.generation.megatron.utils.log_gpu_memory(tag: str) → None#: Print a one-line GPU-memory summary for the calling rank.