diff --git a/model/src/train.py b/model/src/utils/train.py
similarity index 95%
rename from model/src/train.py
rename to model/src/utils/train.py
index 20c7370..4d84f24 100644
--- a/model/src/train.py
+++ b/model/src/utils/train.py
@@ -6,7 +6,7 @@ import torch
 import torch.nn as nn
 from dataloader import VCDADataset  # noqa: F401
 
-from .models import VisualCommunicationModel
+from ..models import VisualCommunicationModel
 
 PRE_WARMUP_LR = 1e-10
 POST_WARMUP_LR = 1e-5