reference/generated/torchrl.reset_cuda_peak_stats

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

reset_cuda_peak_stats

class torchrl.reset_cuda_peak_stats(device: device | int | str | None = None)

Reset the peak-memory counters for device.

Thin wrapper around torch.cuda.reset_peak_memory_stats(). No-op when CUDA is unavailable or device is non-CUDA.

Parameters:

device – CUDA device whose peaks should be cleared. None (default) targets the current CUDA device.

Examples

>>> from torchrl import reset_cuda_peak_stats
>>> reset_cuda_peak_stats()  # safe even without CUDA

See also

cuda_memory_stats(), cuda_memory_profile.