diff --git a/lazyslide/__init__.py b/lazyslide/__init__.py
index 762ad3a..f77b8de 100644
--- a/lazyslide/__init__.py
+++ b/lazyslide/__init__.py
@@ -1,10 +1,7 @@
 """Working with whole slide imaging"""
 __version__ = "0.1.0"
 
-from .wsi import WSI, create_tiles_coords, filter_tiles
-from .readers.utils import get_crop_left_top_width_height
-from .cv_mods import TissueDetectionHE
-from .h5 import H5File
+from .wsi import WSI
 from .utils import get_reader
 
 
diff --git a/lazyslide/h5.py b/lazyslide/h5.py
index ae42d70..4cc78fe 100644
--- a/lazyslide/h5.py
+++ b/lazyslide/h5.py
@@ -2,14 +2,12 @@
 from pathlib import Path
 
 import h5py
-import numpy as np
 
 from .utils import TileOps
 
 
 class H5File:
     """A class that handle storage and reading of h5 file"""
-    # TODO: Add options to save contours and holes
 
     COORDS_KEY = "coords"
     MASKS_KEY = "masks"
diff --git a/lazyslide/torch_dataset.py b/lazyslide/torch_dataset.py
deleted file mode 100644
index 9d1c6a9..0000000
--- a/lazyslide/torch_dataset.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-from torchvision.transforms.v2 import ToImage, ToDtype, Normalize, Compose, Resize
-
-from .normalizer import ColorNormalizer
-
-
-class WSIDataset(Dataset):
-
-    def __init__(self,
-                 wsi,
-                 transform=None,
-                 run_pretrained=False,
-                 color_normalize=None,
-                 ):
-        self.wsi = wsi
-        self.tiles_coords = self.wsi.h5_file.get_coords()
-        self.tile_ops = self.wsi.h5_file.get_tile_ops()
-
-        if run_pretrained:
-            mean = (0.485, 0.456, 0.406)
-            std = (0.229, 0.224, 0.225)
-        else:
-            mean = (0.5, 0.5, 0.5)
-            std = (0.5, 0.5, 0.5)
-
-        self.transform_ops = [ToImage(), ToDtype(dtype=torch.float32),
-                              Normalize(mean=mean, std=std)]
-        if transform is not None:
-            self.transform_ops = [transform]
-        if self.tile_ops.downsample != 1:
-            self.transform_ops = [Resize(size=(self.tile_ops.height,
-                                               self.tile_ops.width))]\
-                                 + self.transform_ops
-        if color_normalize is not None:
-            self.transform_ops = [ColorNormalizer(method=color_normalize)] + self.transform_ops
-        self.transform = Compose(self.transform_ops)
-
-    def __len__(self):
-        return len(self.tiles_coords)
-
-    def __getitem__(self, idx):
-        coords = self.tiles_coords[idx]
-        x, y = coords
-        image = self.wsi.get_patch(y, x,
-                                   self.tile_ops.ops_width,
-                                   self.tile_ops.ops_height,
-                                   self.tile_ops.level)
-        return self.transform(image)
diff --git a/lazyslide/utils.py b/lazyslide/utils.py
index cd6e13f..5c165bb 100644
--- a/lazyslide/utils.py
+++ b/lazyslide/utils.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from itertools import tee
 from typing import Type
 
 from .readers.base import ReaderBase
@@ -6,6 +7,13 @@
 from .readers.openslide import OpenSlideReader
 
 
+def pairwise(iterable):
+    # pairwise('ABCDEFG') --> AB BC CD DE EF FG
+    a, b = tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
 def get_reader(reader="auto") -> Type[ReaderBase]:
     """Return an available backend"""
 
diff --git a/lazyslide/wsi.py b/lazyslide/wsi.py
index 34337a6..d5c4ea2 100644
--- a/lazyslide/wsi.py
+++ b/lazyslide/wsi.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import warnings
-from itertools import tee
 from numbers import Integral
 from pathlib import Path
 from typing import Iterable
@@ -11,20 +10,12 @@
 from matplotlib import pyplot as plt
 from matplotlib.collections import PatchCollection
 from matplotlib.patches import Rectangle
-from numba import njit, prange
+from numba import njit
 
 from .cv_mods import TissueDetectionHE
 from .h5 import H5File
 from .readers.base import ReaderBase
-from .torch_dataset import WSIDataset
-from .utils import get_reader, TileOps
-
-
-def pairwise(iterable):
-    # pairwise('ABCDEFG') --> AB BC CD DE EF FG
-    a, b = tee(iterable)
-    next(b, None)
-    return zip(a, b)
+from .utils import pairwise, get_reader, TileOps
 
 
 @njit
@@ -576,9 +567,8 @@ def plot_mask(self,
             fig.savefig(savefig, **save_kws)
         return ax
 
-    def to_dataset(self, transform=None, run_pretrained=False, **kwargs):
-        # TODO: Allow resize transform on-the-fly to fit into different models
-        return WSIDataset(self, transform=transform, run_pretrained=run_pretrained, **kwargs)
+    # def to_dataset(self, transform=None, run_pretrained=False, **kwargs):
+    #     return WSIDataset(self, transform=transform, run_pretrained=run_pretrained, **kwargs)
 
     def get_patch(self, left, top, width, height, level=0, **kwargs):
         return self.reader.get_patch(left, top, width, height, level=level, **kwargs)
diff --git a/pyproject.toml b/pyproject.toml
index 2ed740e..96a918f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,9 @@ dependencies = [
 vips = [
     "pyvips",
 ]
+all = [
+    "ncls"
+]
 dev = [
     "flit",
     "pytest >=2.7.3",