diff --git a/src/diart/functional.py b/src/diart/functional.py
index 1d340c74..04ba8d50 100644
--- a/src/diart/functional.py
+++ b/src/diart/functional.py
@@ -1,4 +1,6 @@
-from typing import Union, Optional, List, Literal, Iterable, Tuple
+from typing import Union, Optional, List, Iterable, Tuple
+from typing_extensions import Literal 
+
 
 import numpy as np
 import torch
diff --git a/src/diart/sinks.py b/src/diart/sinks.py
index 0f54f8a1..7d0bb8d2 100644
--- a/src/diart/sinks.py
+++ b/src/diart/sinks.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 from traceback import print_exc
-from typing import Literal, Union, Text, Optional, Tuple
+from typing import Union, Text, Optional, Tuple
+from typing_extensions import Literal
 
 import matplotlib.pyplot as plt
 from pyannote.core import Annotation, Segment, SlidingWindowFeature, notebook