From e563e1bf37182290866a1a1b1f5a83e8bf56efd5 Mon Sep 17 00:00:00 2001 From: Chintan Shah Date: Sun, 6 Oct 2019 14:13:02 -0400 Subject: [PATCH] moving tensors to GPU [v3] --- model/pytorch/dcrnn_cell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py index 2a0b85f..7550501 100644 --- a/model/pytorch/dcrnn_cell.py +++ b/model/pytorch/dcrnn_cell.py @@ -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)