Fix prepare_adj_mat() by specifying shape explicitly.

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