dgl.sparse.SparseMatrix.cuda

SparseMatrix.cuda()[源代码]

将矩阵移动到 GPU。如果矩阵已在 GPU 上,则返回原始矩阵。如果存在多个 GPU 设备,将选择 cuda:0

返回:

位于 GPU 上的矩阵

返回类型:

SparseMatrix

示例

>>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]])
>>> A = dglsp.spmatrix(indices, shape=(3, 4))
>>> A.cuda()
SparseMatrix(indices=tensor([[1, 1, 2],
                             [1, 2, 0]], device='cuda:0'),
             values=tensor([1., 1., 1.], device='cuda:0'),
             shape=(3, 4), nnz=3)