diff --git a/datasetss/S3DIS.py b/datasetss/S3DIS.py
index 83371ed..5270d3d 100644
--- a/datasetss/S3DIS.py
+++ b/datasetss/S3DIS.py
@@ -30,9 +30,9 @@ import torch
 import warnings
 from multiprocessing import Lock
 
-
+from utils.ply import write_ply, read_ply
 # OS functions
-from os import listdir
+from os import listdir, makedirs
 from os.path import exists, join, isdir
 
 # Dataset parent class
@@ -757,6 +757,7 @@ class S3DISDataset(PointCloudDataset):
                         cloud_classes = np.vstack((cloud_classes, object_classes))
 
             # Save as ply
+
             write_ply(
                 cloud_file,
                 (cloud_points, cloud_colors, cloud_classes),
@@ -830,6 +831,8 @@ class S3DISDataset(PointCloudDataset):
                 sub_labels = np.squeeze(sub_labels)
 
                 # Get chosen neighborhoods
+                from sklearn.neighbors import KDTree
+
                 search_tree = KDTree(sub_points, leaf_size=10)
                 # search_tree = nnfln.KDTree(n_neighbors=1, metric='L2', leaf_size=10)
                 # search_tree.fit(sub_points)