class SkipBlock(nn.Module):
    def __init__(self, dim: int):
        super().__init__()

        self.skip_linear = nn.Linear(2 * dim, dim)

        # Use torch.nn.LayerNorm for now, following the original code
        self.norm = nn.LayerNorm(dim)

    def forward(self, x, skip):
        x = self.skip_linear(torch.cat([x, skip], dim=-1))
        x = self.norm(x)

        return x
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/unidiffuser/modeling_uvit.py