changed to cuda loading for Aquila base model

ftgreat · ftgreat · commit 37e7cf7f33b7 · 2023-07-05T17:57:33.000+08:00
Signed-off-by: ftgreat &lt;ftgreat@163.com&gt;
diff --git a/examples/Aquila/Aquila-pretrain/generate.py b/examples/Aquila/Aquila-pretrain/generate.py
@@ -12,6 +12,7 @@
                     model_dir=state_dict,
                     model_name=model_name,
                     use_cache=True,
+                    device='cuda',
                     fp16=True)
 model = loader.get_model()
 tokenizer = loader.get_tokenizer()
diff --git a/examples/Aquila/generate.py b/examples/Aquila/generate.py
@@ -14,6 +14,7 @@
 loader = AutoLoader("lm",
                     model_dir=state_dict,
                     model_name=model_name,
+                    device='cuda',
                     use_cache=True,
                     fp16=True)
 model = loader.get_model()
