Update mplug_docowl/model/builder.py

mplug_docowl/model/builder.py

@@ -39,8 +39,8 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
             bnb_4bit_quant_type='nf4'
         )
     else:
-        # kwargs['torch_dtype'] = torch.float16
-        kwargs['torch_dtype'] = torch.bfloat16
+        kwargs['torch_dtype'] = torch.float16
+
     if 'paperowl' or 'docowl' in model_name.lower():
         if model_base is not None:
             # this may be mm projector only
@@ -48,6 +48,8 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
             cfg_pretrained = AutoConfig.from_pretrained(model_path)
             model = MPLUGDocOwlLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+            # adjust for zero environment of huggingface space
+            model.to("cuda:0")
         else:
             tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
             model = MPLUGDocOwlLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)