Update models.py

models.py

@@ -18,7 +18,7 @@ class Models():
         #is it best to keep in memory why not pickle?
         checkpoint = "Universal-NER/UniNER-7B-all"
         
-        ner_model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16, offload_folder="offload", offload_state_dict = True)
+        ner_model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float32, offload_folder="offload", offload_state_dict = True)
         tokenizer = AutoTokenizer.from_pretrained("Universal-NER/UniNER-7B-all", use_fast=False, padding="max_length")
         hf_pipeline = pipeline(
             "text-generation", #task