Fix import paths in pretraining scripts

Luodian · Nov 6, 2023 · ca69589 · ca69589
1 parent b1557c7
commit ca69589
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pipeline/demos/demo_models.py b/pipeline/demos/demo_models.py
@@ -166,7 +166,7 @@ def generate(self, image, prompt, no_image_flag=False):
         formated_prompt = f"User: {prompt} Assistant:"
         model_inputs = self.processor(text=formated_prompt, images=[raw_image_data] if no_image_flag is False else None, device=self.device)
         for k, v in model_inputs.items():
-            model_inputs[k] = v.to(self.device)
+            model_inputs[k] = v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else [vv.to(self.device, non_blocking=True) for vv in v]
 
         generation_output = self.model.generate(**model_inputs, max_new_tokens=self.max_new_tokens)
         generation_text = self.processor.batch_decode(generation_output, skip_special_tokens=True)

diff --git a/pipeline/train/pretraining.py b/pipeline/train/pretraining.py
@@ -23,7 +23,7 @@
 from otter_ai import FlamingoForConditionalGeneration, OtterForConditionalGeneration
 
 sys.path.append("../..")
-from pipeline.train.data import get_data
+from pipeline.mimicit_utils.data import get_data
 from pipeline.train.distributed import world_info_from_env
 from pipeline.train.train_utils import AverageMeter, get_checkpoint
 

diff --git a/pipeline/train/pretraining_cc3m.py b/pipeline/train/pretraining_cc3m.py
@@ -23,7 +23,7 @@
 from otter_ai import FlamingoForConditionalGeneration, OtterForConditionalGeneration
 
 sys.path.append("../..")
-from pipeline.train.data import get_data
+from pipeline.mimicit_utils.data import get_data
 from pipeline.train.distributed import world_info_from_env
 from pipeline.train.train_utils import AverageMeter, get_checkpoint