def gib(value):
    return value / 1024 ** 3

try:
    train_one_epoch(model, loader, optimizer)
except torch.cuda.OutOfMemoryError:
    print(f"allocated={gib(torch.cuda.memory_allocated()):.2f} GiB")
    print(f"reserved={gib(torch.cuda.memory_reserved()):.2f} GiB")
    print(f"peak_allocated={gib(torch.cuda.max_memory_allocated()):.2f} GiB")
    print(torch.cuda.memory_summary(abbreviated=True))
    raise