souging/TRELLIS_TextTo3D

1import torch2import torch.nn as nn3from . import SparseTensor4 5__all__ = [6    'SparseReLU',7    'SparseSiLU',8    'SparseGELU',9    'SparseActivation'10]11 12 13class SparseReLU(nn.ReLU):14    def forward(self, input: SparseTensor) -> SparseTensor:15        return input.replace(super().forward(input.feats))16    17 18class SparseSiLU(nn.SiLU):19    def forward(self, input: SparseTensor) -> SparseTensor:20        return input.replace(super().forward(input.feats))21 22 23class SparseGELU(nn.GELU):24    def forward(self, input: SparseTensor) -> SparseTensor:25        return input.replace(super().forward(input.feats))26 27 28class SparseActivation(nn.Module):29    def __init__(self, activation: nn.Module):30        super().__init__()31        self.activation = activation32 33    def forward(self, input: SparseTensor) -> SparseTensor:34        return input.replace(self.activation(input.feats))35    36