Function print_peak_memory

diffportrait360_release/code/utils/utils.py:36–38

(prefix, device)

    
# display the peak memory of cuda
def print_peak_memory(prefix, device):
    if device == 0:
        print(f"{prefix}: {torch.cuda.max_memory_allocated(device) // 1e6}MB ")

def anal_tensor(tensor, name):
    sent = f" name: {name} mean: {tensor.mean().item()}  std: {tensor.std().item()}  min: {tensor.min().item()}  max: {tensor.max().item()}"