diff --git a/codes/models/image_latents/vit_latent.py b/codes/models/image_latents/vit_latent.py
index e988cf1f..0243e8a7 100644
--- a/codes/models/image_latents/vit_latent.py
+++ b/codes/models/image_latents/vit_latent.py
@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from codes.models.arch_util import ResBlock
-from codes.models.lucidrains.x_transformers import Encoder
-from codes.trainer.networks import register_model
+from models.arch_util import ResBlock
+from models.lucidrains.x_transformers import Encoder
+from trainer.networks import register_model
 
 
 class VitLatent(nn.Module):