diff --git a/Pretrain_nlvr.py b/Pretrain_nlvr.py
index 8fb5a1c..c05bf88 100644
--- a/Pretrain_nlvr.py
+++ b/Pretrain_nlvr.py
@@ -15,7 +15,7 @@
 import torch.backends.cudnn as cudnn
 import torch.distributed as dist
 
-from models.model_retrieval import ALBEF
+from models.model_pretrain_nlvr import ALBEF
 from models.vit import interpolate_pos_embed
 from models.tokenization_bert import BertTokenizer
 
@@ -197,4 +197,4 @@ def main(args, config):
 
     yaml.dump(config, open(os.path.join(args.output_dir, 'config.yaml'), 'w'))    
     
-    main(args, config)
\ No newline at end of file
+    main(args, config)