diff --git a/dask/diagnostics/cache.py b/dask/diagnostics/cache.py
index e2c65a6fc3fc..9a38ea8bd57e 100644
--- a/dask/diagnostics/cache.py
+++ b/dask/diagnostics/cache.py
@@ -1,4 +1,4 @@
-from .core import Diagnostic
+from ..callbacks import Callback
 from timeit import default_timer
 from numbers import Number
 
@@ -7,7 +7,7 @@
 except ImportError:
     pass
 
-class Cache(Diagnostic):
+class Cache(Callback):
     """ Use cache for computation
 
     Example