diff --git a/image_prediction/label_mapper/mappers/probability.py b/image_prediction/label_mapper/mappers/probability.py
index cb2d3f8..9e7347c 100644
--- a/image_prediction/label_mapper/mappers/probability.py
+++ b/image_prediction/label_mapper/mappers/probability.py
@@ -1,4 +1,3 @@
-from functools import partial
 from operator import itemgetter
 from typing import Mapping, Iterable