I try to create a multi-dimensional sparse tensor in my implemention of sparse GAT, but I get the following error.
class SpecialSpmmFunction(torch.autograd.Function):
    """Special function for only sparse region backpropataion layer."""
    @staticmethod
    def forward(ctx, indices, values, shape, b):
        assert indices.requires_grad == False
        a = torch.sparse_coo_tensor(indices, values, shape) #  This error occurred
        ctx.save_for_backward(a, b)
        ctx.N = shape[0]
        return torch.matmul(a, b)

    @staticmethod
    def backward(ctx, grad_output):
       # ............
class SpecialSpmm(nn.Module):
    def forward(self, indices, values, shape, b):
        return SpecialSpmmFunction.apply(indices, values, shape, b)
#.............
 e_rowsum = self.special_spmm(edge.unsqueeze(0).expand(nbatches, -1, -1), edge_e, torch.Size([nbatches, N, N]), torch.ones(size=(N, 1), device=dv))


a = torch.sparse_coo_tensor(indices, values, shape)
RuntimeError: indices must be sparse_dim x nnz, but got: [161, 2, 3025]

torch.sparse_coo_tensor seems like that cannot create multi-dimensional sparse tensor