diff --git a/train/model.py b/train/model.py
index 1e4d4af..15b2f1b 100644
--- a/train/model.py
+++ b/train/model.py
@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 from dataclasses import dataclass
 
 @dataclass