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): def _build_sparse_matrix(L):
L = L.tocoo() L = L.tocoo()
indices = np.column_stack((L.row, L.col)) 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 L
# return torch.sparse.sparse_reorder(L) # return torch.sparse.sparse_reorder(L)