Update Pretrain_nlvr.py

SuJQ0722 · Oct 24, 2021 · 245a86a · 245a86a
1 parent cd93adf
commit 245a86a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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)
+    main(args, config)