diff --git a/eigen_attention.py b/eigen_attention.py
index 1c06945..6148ddb 100644
--- a/eigen_attention.py
+++ b/eigen_attention.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
diff --git a/eigen_memory.py b/eigen_memory.py
index 9a9b7dd..d043c53 100644
--- a/eigen_memory.py
+++ b/eigen_memory.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
diff --git a/eigen_transformer.py b/eigen_transformer.py
index 0ecc6bf..22d8dca 100644
--- a/eigen_transformer.py
+++ b/eigen_transformer.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 import torch.nn as nn