diff --git a/src/icosagon/normalize.py b/src/icosagon/normalize.py index ec09b3c..e13fb05 100644 --- a/src/icosagon/normalize.py +++ b/src/icosagon/normalize.py @@ -41,7 +41,8 @@ def _sparse_coo_tensor(indices, values, size): torch.uint8: torch.sparse.ByteTensor, torch.long: torch.sparse.LongTensor, torch.int: torch.sparse.IntTensor, - torch.short: torch.sparse.ShortTensor }[values.dtype] + torch.short: torch.sparse.ShortTensor, + torch.bool: torch.sparse.ByteTensor }[values.dtype] return ctor(indices, values, size) diff --git a/tests/icosagon/test_normalize.py b/tests/icosagon/test_normalize.py index e4b0fef..f3de835 100644 --- a/tests/icosagon/test_normalize.py +++ b/tests/icosagon/test_normalize.py @@ -46,14 +46,14 @@ def test_add_eye_sparse_04(): def test_norm_adj_mat_one_node_type_sparse_01(): adj_mat = torch.rand((10, 10)) - adj_mat = (adj_mat > .5) + adj_mat = (adj_mat > .5).to(torch.float32) adj_mat = adj_mat.to_sparse() _ = norm_adj_mat_one_node_type_sparse(adj_mat) def test_norm_adj_mat_one_node_type_sparse_02(): adj_mat_dense = torch.rand((10, 10)) - adj_mat_dense = (adj_mat_dense > .5) + adj_mat_dense = (adj_mat_dense > .5).to(torch.float32) adj_mat_sparse = adj_mat_dense.to_sparse() adj_mat_sparse = norm_adj_mat_one_node_type_sparse(adj_mat_sparse) adj_mat_dense = norm_adj_mat_one_node_type_dense(adj_mat_dense) diff --git a/tests/icosagon/test_trainprep.py b/tests/icosagon/test_trainprep.py index b49646b..5ec86a7 100644 --- a/tests/icosagon/test_trainprep.py +++ b/tests/icosagon/test_trainprep.py @@ -108,7 +108,7 @@ def test_prepare_adj_mat_02(): def test_prepare_relation_type_01(): - adj_mat = (torch.rand((10, 10)) > .5) + adj_mat = (torch.rand((10, 10)) > .5).to(torch.float32) r = RelationType('Test', 0, 0, adj_mat, True) ratios = TrainValTest(.8, .1, .1) _ = prepare_relation_type(r, ratios, False)