dgl.sparse.SparseMatrix.cuda
- SparseMatrix.cuda()[源代码]
将矩阵移动到 GPU。如果矩阵已在 GPU 上,则返回原始矩阵。如果存在多个 GPU 设备,将选择
cuda:0
。- 返回:
位于 GPU 上的矩阵
- 返回类型:
示例
>>> 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)