czzhangheng
/
DCRNN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

moving tensors to GPU [v3]

This commit is contained in:
Chintan Shah 2019-10-06 14:13:02 -04:00
parent 017ec70783
commit e563e1bf37
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model/pytorch/dcrnn_cell.py
View File

@ -74,7 +74,7 @@ class DCGRUCell(torch.nn.Module):
def _build_sparse_matrix(L):
L = L.tocoo()
indices = np.column_stack((L.row, L.col))
L = torch.sparse_coo_tensor(indices.T, L.data, L.shape)
L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device)
return L
# return torch.sparse.sparse_reorder(L)