Update hf-inference-example.py

3 years ago · 233196eda1
parent 8a484117f5
commit 233196eda1
1 changed files with 4 additions and 0 deletions
--- a/hf-inference-example.py
+++ b/hf-inference-example.py
@ -1,5 +1,9 @@
 import llamahf

+# to save memory use bfloat16 on cpu
+# import torch
+# torch.set_default_dtype(torch.bfloat16)
+
 MODEL = 'decapoda-research/llama-7b-hf'
 # MODEL = 'decapoda-research/llama-13b-hf'
 # MODEL = 'decapoda-research/llama-30b-hf'