Update open_flamingo/open_flamingo/src/factory.py

open_flamingo/open_flamingo/src/factory.py

@@ -15,6 +15,7 @@ def create_model_and_transforms(
     use_local_files: bool = False,
     decoder_layers_attr_name: str = None,
     freeze_lm_embeddings: bool = False,
+    device: int = 0,
     **flamingo_kwargs,
 ):
     """
@@ -39,6 +40,7 @@ def create_model_and_transforms(
     )
     # set the vision encoder to output the visual features
     vision_encoder.visual.output_tokens = True
+        vision_encoder.to(device, dtype=torch.bfloat16) if device > -1 else None
 
     text_tokenizer = AutoTokenizer.from_pretrained(
         tokenizer_path,
@@ -58,7 +60,8 @@ def create_model_and_transforms(
         lang_encoder_path,
         local_files_only=use_local_files,
         trust_remote_code=True,
-    )
+        
+    ).to(device, dtype=torch.bfloat16) if device > -1 else None
 
     # hacks for MPT-1B, which doesn't have a get_input_embeddings method
     if "mpt-1b-redpajama-200b" in lang_encoder_path:
@@ -90,7 +93,7 @@ def create_model_and_transforms(
         ],
         cross_attn_every_n_layers=cross_attn_every_n_layers,
         **flamingo_kwargs,
-    )
+    ).to(device, dtype=torch.bfloat16) if device > -1 else None
 
     # Freeze all parameters
     model.requires_grad_(False)