dsmm_test.py 文件源码

def test_forward_batch():
    i = torch.LongTensor([[0, 0, 0, 1, 1, 1],
                          [0, 1, 1, 0, 1, 1],
                          [2, 0, 2, 2, 0, 2]])
    v = torch.FloatTensor([3, 4, 5, 6, 7, 8])
    sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 3]))
    dense = Variable(torch.randn(2, 3, 3))

    res = gpytorch.dsmm(Variable(sparse), dense)
    actual = torch.matmul(Variable(sparse.to_dense()), dense)
    assert(torch.norm(res.data - actual.data) < 1e-5)