diff --git a/pointcnn/core.py b/pointcnn/core.py
index d852f54..f4fe334 100644
--- a/pointcnn/core.py
+++ b/pointcnn/core.py
@@ -140,7 +140,7 @@ class PointCNN(nn.Module):
         return regions
 
     # @timed.timed
-    def forward(self, ps, P, F):
+    def forward(self, ps, P, F, P_idx = None):
         """
         Given a set of representative points, a point cloud, and its
         corresponding features, return a new set of representative points with
@@ -154,7 +154,8 @@ class PointCNN(nn.Module):
         :param F: Regional features such that P[:,p_idx,:] is the feature associated with F[:,p_idx,:]
         :return:
         """
-        P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors).cuda()  # This step takes ~97% of the time.
+        if P_idx != None:
+            P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors).cuda()  # This step takes ~97% of the time.
         P_regional = self.select_region(P, P_idx)                           # Prime target for optimization: KNN on GPU.
         if False:
             # Draw neighborhood points, for debugging.
diff --git a/pointcnn/util.py b/pointcnn/util.py
index f31a30c..01c538a 100644
--- a/pointcnn/util.py
+++ b/pointcnn/util.py
@@ -3,6 +3,8 @@ import torch
 import numpy as np
 from sklearn.neighbors import NearestNeighbors
 
+torch.CUDA_LAUNCH_BLOCKING = 1
+
 try:
     from .context import pytorch_knn_cuda
 except SystemError:
@@ -67,7 +69,6 @@ def knn_indices_func_gpu(ps, P, k):
 
     def single_batch_knn(p, P_particular):
         nbrs_f = pytorch_knn_cuda.KNearestNeighbor(k + 1)
-        # knn_cuda(k + 1, )
         indices = nbrs_f(P_particular, p)[0]
         return indices[:,1:]
 
@@ -76,12 +77,11 @@ def knn_indices_func_gpu(ps, P, k):
     ], dim = 0)
     return region_idx
 
-
 if __name__ == "__main__":
     from torch.autograd import Variable
-    N_rep = 100
+    N_rep = 1000
     N = 2
-    num_points = 1000
+    num_points = 10000
     D = 3
     test_P  = np.random.rand(N,num_points,D).astype(np.float32)
     idx = np.random.choice(test_P.shape[1], N_rep, replace = False)