diff --git a/ciw/dists/distributions.py b/ciw/dists/distributions.py
index 31be896..17030ec 100644
--- a/ciw/dists/distributions.py
+++ b/ciw/dists/distributions.py
@@ -11,7 +11,11 @@
     lognormvariate,
     weibullvariate,
 )
-from typing import List, NoReturn
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import List, NoReturn
 
 import numpy as np