diff --git a/edgeconv.py b/edgeconv.py
index ac9bd212eb3e24aafcdbfa2af9775478c6701121..0686f1dcb3154d509a12a985f94b6b0e9b14c414 100644
--- a/edgeconv.py
+++ b/edgeconv.py
@@ -147,5 +147,5 @@ def knn(topk_indices, features):
     # features: (N, P, C)
     # return: (N, P, K, C)
     with tf.name_scope("knn"):
-        knn_features = tf.gather(features, topk_indices, batch_dims=1)
-    return knn_features
+        topk_features = tf.gather(features, topk_indices, batch_dims=1)
+        return topk_features