Spaces:
Running on Zero
Running on Zero
File size: 522 Bytes
afea36f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | # Copied from the TRELLIS project:
# https://github.com/microsoft/TRELLIS
# Original license: MIT
# Copyright (c) the TRELLIS authors
import torch
import torch.nn as nn
from . import SparseTensor
__all__ = [
'SparseLinear'
]
class SparseLinear(nn.Linear):
def __init__(self, in_features, out_features, bias=True):
super(SparseLinear, self).__init__(in_features, out_features, bias)
def forward(self, input: SparseTensor) -> SparseTensor:
return input.replace(super().forward(input.feats))
|