[update]edit main

main.py

@@ -64,7 +64,7 @@ def main():
         args.pretrained_model_name_or_path,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.bfloat16,
         device_map="auto",
         offload_folder="./offload",
         offload_state_dict=True,
@@ -78,6 +78,7 @@ def main():
             return_tensors="pt",
             add_special_tokens=False,
         ).input_ids.to(args.device)
+
         with torch.no_grad():
             outputs = model.generate(
                 input_ids=input_ids, max_new_tokens=args.max_new_tokens, do_sample=True,