I think it would be nicer if SparseLinear was just a part of the Linear module. The missing functionality for sparse tensors to work with the Linear model should just be the matmul function, which is here: https://github.com/pytorch/pytorch/blob/30ec06c140b0428d591e2f5007bc8046d1bdf7c4/aten/src/ATen/native/LinearAlgebra.cpp#L142-L161 14
You would need to implement sparse * sparse matrix multiplication. There are some libraries out there that do it, but in general, I think the easiest way to get this to work is by converting the sparse matrices to CSR format (pytorch sparse matrices are COO format) and performing the sparse multiply.