From 0578b13ea4c04aa8b93e0fc31d4a89293f8c7b9d Mon Sep 17 00:00:00 2001 From: Stanislaw Adaszewski Date: Wed, 17 Jun 2020 12:14:05 +0200 Subject: [PATCH] Fix prepare_adj_mat() by specifying shape explicitly. --- src/icosagon/trainprep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/icosagon/trainprep.py b/src/icosagon/trainprep.py index 3457791..6e8211b 100644 --- a/src/icosagon/trainprep.py +++ b/src/icosagon/trainprep.py @@ -111,7 +111,7 @@ def prepare_adj_mat(adj_mat: torch.Tensor, edges_neg = train_val_test_split_edges(edges_neg, ratios) adj_mat_train = torch.sparse_coo_tensor(indices = edges_pos.train.transpose(0, 1), - values=torch.ones(len(edges_pos.train), dtype=adj_mat.dtype)) + values=torch.ones(len(edges_pos.train), shape=adj_mat.shape, dtype=adj_mat.dtype)) return adj_mat_train, edges_pos, edges_neg