File size: 387 Bytes
394918c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | 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))
|