Function count_param

diffportrait360_release/code/utils/utils.py:32–33

(module_list)

# count # of param for a list of module
def count_param(module_list):
    return sum(x.numel() for module in module_list for x in module.parameters()) / 10**6
    
# display the peak memory of cuda
def print_peak_memory(prefix, device):