فهرست منبع

Use the CPU if no GPU is detected

oobabooga 3 سال پیش
والد
کامیت
5d3f15b915
1فایلهای تغییر یافته به همراه3 افزوده شده و 0 حذف شده
  1. 3 0
      server.py

+ 3 - 0
server.py

@@ -120,6 +120,9 @@ def load_model(model_name):
     else:
         command = "AutoModelForCausalLM.from_pretrained"
         params = ["low_cpu_mem_usage=True"]
+        if not args.cpu and not torch.cuda.is_available():
+            print("Warning: no GPU has been detected.\nFalling back to CPU mode.\n")
+            args.cpu = True
 
         if args.cpu:
             params.append("low_cpu_mem_usage=True")