Commit 0fb1f1db authored by rusty1s's avatar rusty1s
Browse files

version up

parent 4fa1fb22
#include <ATen/CPUGenerator.h>
#include <torch/extension.h> #include <torch/extension.h>
at::Tensor neighbor_sampler(at::Tensor start, at::Tensor cumdeg, size_t size, at::Tensor neighbor_sampler(at::Tensor start, at::Tensor cumdeg, size_t size,
float factor) { float factor) {
at::CPUGenerator *generator = at::detail::getDefaultCPUGenerator();
auto start_ptr = start.data<int64_t>(); auto start_ptr = start.data<int64_t>();
auto cumdeg_ptr = cumdeg.data<int64_t>(); auto cumdeg_ptr = cumdeg.data<int64_t>();
...@@ -24,7 +22,7 @@ at::Tensor neighbor_sampler(at::Tensor start, at::Tensor cumdeg, size_t size, ...@@ -24,7 +22,7 @@ at::Tensor neighbor_sampler(at::Tensor start, at::Tensor cumdeg, size_t size,
std::unordered_set<int64_t> set; std::unordered_set<int64_t> set;
if (size_i < 0.7 * float(num_neighbors)) { if (size_i < 0.7 * float(num_neighbors)) {
while (set.size() < size_i) { while (set.size() < size_i) {
int64_t z = generator->random() % num_neighbors; int64_t z = rand() % num_neighbors;
set.insert(z + low); set.insert(z + low);
} }
std::vector<int64_t> v(set.begin(), set.end()); std::vector<int64_t> v(set.begin(), set.end());
......
...@@ -29,7 +29,7 @@ if CUDA_HOME is not None: ...@@ -29,7 +29,7 @@ if CUDA_HOME is not None:
['cuda/rw.cpp', 'cuda/rw_kernel.cu']), ['cuda/rw.cpp', 'cuda/rw_kernel.cu']),
] ]
__version__ = '1.4.3a1' __version__ = '1.4.3'
url = 'https://github.com/rusty1s/pytorch_cluster' url = 'https://github.com/rusty1s/pytorch_cluster'
install_requires = ['scipy'] install_requires = ['scipy']
......
...@@ -7,7 +7,7 @@ from .radius import radius, radius_graph ...@@ -7,7 +7,7 @@ from .radius import radius, radius_graph
from .rw import random_walk from .rw import random_walk
from .sampler import neighbor_sampler from .sampler import neighbor_sampler
__version__ = '1.4.3a1' __version__ = '1.4.3'
__all__ = [ __all__ = [
'graclus_cluster', 'graclus_cluster',
......
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