Suppose I perform my_awesome_cuda_training() and I want to determine the maximum amount of memory that was used. Is that as simple as using torch.cuda.memory_stats()?
import torch


def my_awesome_cuda_training():
    torch.empty(2, *[1024] * 3, dtype=torch.uint8, device="cuda")


my_awesome_cuda_training()

stats = torch.cuda.memory_stats()
peak_bytes_requirement = stats["allocated_bytes.all.peak"]
print(f"Peak memory requirement: {peak_bytes_requirement / 1024 ** 3:.2f} GB")

Peak memory requirement: 2.00 GB

It seems to work as intended, but I would be grateful for confirmation.