Starbourne
/

cogvlm-chat-hf

Text Generation

Inference Endpoints

Model card Files Files and versions Community

Marlon Wiprud commited on Nov 28, 2023

Commit

a894d5f

·

1 Parent(s): 41902fa

chore: back to bf16

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -55,7 +55,7 @@ class EndpointHandler:
         with init_empty_weights():
             self.model = AutoModelForCausalLM.from_pretrained(
                 "THUDM/cogvlm-chat-hf",
-                torch_dtype=torch.float16,
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
             )

         with init_empty_weights():
             self.model = AutoModelForCausalLM.from_pretrained(
                 "THUDM/cogvlm-chat-hf",
+                torch_dtype=torch.bfloat16,
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
             )