smarter cumDegree

aten/THC/THCDegree.cuh

@@ -8,6 +8,7 @@ template<typename T>
 __global__ void cumDegreeKernel(T *self, int64_t *index, ptrdiff_t nEdges) {
   KERNEL_LOOP(i, nEdges) {
     int64_t r = index[i];
+    if (i + 1 == nEdges) {self[r] = ScalarConvert<int, T>::to(nEdges);; continue; }
     if (r != index[i+1]) { self[r] = ScalarConvert<int, T>::to(i + 1); }
   }
 }

aten/THC/THCGraclus.cu

@@ -20,7 +20,7 @@ void THCTensor_graclus(THCState *state, THCudaLongTensor *self, THCudaLongTensor
   THCudaLongTensor_degree(state, degree, row);
 
   THCudaLongTensor *cumDegree = THCudaLongTensor_newWithSize1d(state, nNodes);
-  THCudaLongTensor_cumDegree(state, cumDegree, row);
+  THCudaLongTensor_cumDegree(state, self, row);
 
   while(!THCTensor_color(state, self)) {
     THCTensor_propose(state, self, prop, row, col, degree, cumDegree);

aten/THC/generic/THCDegree.cuh

@@ -20,7 +20,7 @@ void THCTensor_(cumDegree)(THCState *state, THCTensor *self, THCudaLongTensor *i
   real *selfData = THCTensor_(data)(state, self);
   int64_t *indexData = THCudaLongTensor_data(state, index);
 
-  KERNEL_RUN(cumDegreeKernel, nEdges - 1, selfData, indexData);
+  KERNEL_RUN(cumDegreeKernel, nEdges, selfData, indexData);
 }
 
 #endif  // THC_GENERIC_FILE