Skip to content

GitLab

"tests/python/common/test_partition.py" did not exist on "ed8e9c4440c204d19a138bcc552ee476c3a7cbc1"
Commit 5b9304b9 authored by bowendeng's avatar bowendeng
Browse files

code clean

parent 0090f4ed
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
torch_sparse/metis.py
View file @ 5b9304b9
...@@ -18,9 +18,8 @@ def metis_wgt(x): ...@@ -18,9 +18,8 @@ def metis_wgt(x):
return (x_ratio * arange + tick).long(), tick, arange return (x_ratio * arange + tick).long(), tick, arange
def partition( def partition(src: SparseTensor, num_parts: int, recursive: bool = False
src: SparseTensor, num_parts: int, recursive: bool = False ) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
rowptr, col = src.storage.rowptr().cpu(), src.storage.col().cpu() rowptr, col = src.storage.rowptr().cpu(), src.storage.col().cpu()
edge_wgt = src.storage.value().cpu() edge_wgt = src.storage.value().cpu()
edge_wgt = metis_wgt(edge_wgt)[0] edge_wgt = metis_wgt(edge_wgt)[0]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment