graclus weight gpu bugfix

rusty1s · Apr 18, 2018 · 13dabd4 · 13dabd4
1 parent 7985cdd
commit 13dabd4
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 6 deletions.
diff --git a/aten/THC/THCNumerics.cuh b/aten/THC/THCNumerics.cuh
@@ -16,14 +16,14 @@
 template<typename T>
 struct THCNumerics {
   static inline __host__ __device__ T div(T a, T b) { return a / b; }
-  static inline __host__ __device__ bool gt(T a, T b) { return a > b; }
+  static inline __host__ __device__ bool gte(T a, T b) { return a >= b; }
 };
 
 #ifdef CUDA_HALF_TENSOR
 template<>
 struct THCNumerics<half> {
   static inline __host__ __device__ half div(half a, half b) { return f2h(h2f(a) / h2f(b)); }
-  static inline __host__ __device__ bool gt(half a, half b) { return h2f(a) > h2f(b); }
+  static inline __host__ __device__ bool gte(half a, half b) { return h2f(a) >= h2f(b); }
 };
 #endif  // CUDA_HALF_TENSOR
 

diff --git a/aten/THC/THCPropose.cuh b/aten/THC/THCPropose.cuh
@@ -32,7 +32,7 @@ __global__ void weightedProposeKernel(int64_t *color, int64_t *prop, int64_t *ro
       tmp = weight[e];
       if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
       // Find maximum weighted red neighbor.
-      if (color[c] == -2 && THCNumerics<T>::gt(tmp, maxWeight)) {
+      if (color[c] == -2 && THCNumerics<T>::gte(tmp, maxWeight)) {
         matchedValue = c;
         maxWeight = tmp;
       }

diff --git a/aten/THC/THCResponse.cuh b/aten/THC/THCResponse.cuh
@@ -35,7 +35,7 @@ __global__ void weightedResponseKernel(int64_t *color, int64_t *prop, int64_t *r
       tmp = weight[e];
       if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
       // Find maximum weighted blue neighbor, who proposed to i.
-      if (color[c] == -1 && prop[c] == i && THCNumerics<T>::gt(tmp, maxWeight)) {
+      if (color[c] == -1 && prop[c] == i && THCNumerics<T>::gte(tmp, maxWeight)) {
         matchedValue = c;
         maxWeight = tmp;
       }

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '1.0.2'
+__version__ = '1.0.3'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi']

diff --git a/torch_cluster/__init__.py b/torch_cluster/__init__.py
@@ -1,6 +1,6 @@
 from .graclus import graclus_cluster
 from .grid import grid_cluster
 
-__version__ = '1.0.2'
+__version__ = '1.0.3'
 
 __all__ = ['graclus_cluster', 'grid_cluster', '__version__']