Update open_flamingo/open_flamingo/src/factory.py

open_flamingo/open_flamingo/src/factory.py

@@ -16,7 +16,6 @@ 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,
 ):
     """
@@ -41,7 +40,6 @@ 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,
@@ -60,7 +58,7 @@ def create_model_and_transforms(
     lang_encoder = AutoModelForCausalLM.from_pretrained(
         lang_encoder_path,
         local_files_only=use_local_files,
-        trust_remote_code=True).to(device, dtype=torch.bfloat16) if device > -1 else None
+        trust_remote_code=True)
 
     # hacks for MPT-1B, which doesn't have a get_input_embeddings method
     if "mpt-1b-redpajama-200b" in lang_encoder_path:
@@ -91,7 +89,7 @@ def create_model_and_transforms(
             "width"
         ],
         cross_attn_every_n_layers=cross_attn_every_n_layers,
-        **flamingo_kwargs).to(device, dtype=torch.bfloat16) if device > -1 else None
+        **flamingo_kwargs)
 
     # Freeze all parameters
     model.requires_grad_(False)