Method forward

trellis2/modules/sparse/norm.py:60–64

(self, x: VarLenTensor)

    A LayerNorm layer that converts to float32 before the forward pass.
    """
    def forward(self, x: VarLenTensor) -> VarLenTensor:
        x_dtype = x.dtype
        x = manual_cast(x, torch.float32)
        o = super().forward(x)
        return manual_cast(o, x_dtype)