utils.py

import torch.nn as nn


def select_norm(norm, dim, shape):
    """Just a wrapper to select the normalization type."""
    if norm == "ln":
        return nn.GroupNorm(1, dim, eps=1e-8)