diff --git a/sella/optimize/optimize.py b/sella/optimize/optimize.py
index d33d1ff..414c166 100644
--- a/sella/optimize/optimize.py
+++ b/sella/optimize/optimize.py
@@ -42,7 +42,7 @@ def __init__(
         atoms: Atoms,
         restart: bool = None,
         logfile: str = '-',
-        trajectory: str = None,
+        trajectory: Union[str, Trajectory] = None,
         master: bool = None,
         force_consistent: bool = False,
         delta0: float = None,