Update model.py

src/models/model.py

@@ -8,3 +8,4 @@ from transformers import (
     T5ForConditionalGeneration,
     T5TokenizerFast as T5Tokenizer,
 )
+from torch.utils.data import Dataset, DataLoader