Update aquila_chat.py

SysuJayce · Jun 30, 2023 · e1cc636 · e1cc636
1 parent 893a484
commit e1cc636
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/Aquila/Aquila-chat/aquila_chat.py b/examples/Aquila/Aquila-chat/aquila_chat.py
@@ -83,7 +83,7 @@
 
 config_file = os.path.join(cache_dir, 'config.json')
 from flagai.model.aquila_model import AQUILAModel
-model = AQUILAModel.init_from_json(config_file=config_file)
+model = AQUILAModel.init_from_json(config_file=config_file, device=device, fp16=True) 
 print('*'*20, "model", model)
 
 #lora
@@ -258,4 +258,4 @@ def padding(indice, max_length, pad_idx=0):
     valid_dataset=valid_dataset,
     collate_fn=ConversationDatasetV2.collate_fn,
     optimizer=None,
-    rank_split=False)
+    rank_split=False)