Function count_params

sat/vae_modules/cp_enc_dec.py:717–721

(model, verbose=False)

def count_params(model, verbose=False):
    total_params = sum(p.numel() for p in model.parameters())
    if verbose:
        print(f'{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params')
    return total_params

# * ContextParallelEncoder3D has 92.22 M params.
class ContextParallelEncoder3D(nn.Module):