Spaces:

chansung
/

LLaMA-7B

Runtime error

chansung commited on Mar 3, 2023

Commit

e2ea425

1 Parent(s): 0e25a0e

Update gen.py

Files changed (1) hide show

gen.py CHANGED Viewed

@@ -63,7 +63,7 @@ def get_pretrained_models(
     llama_ckpt_path = checkpoints[local_rank]
     print("Loading")
-    checkpoint = torch.load(llama_ckpt_path, map_location=lambda storage, loc: storage.cuda(0))
     with open(Path(llama_weight_path) / "params.json", "r") as f:
         params = json.loads(f.read())

     llama_ckpt_path = checkpoints[local_rank]
     print("Loading")
+    checkpoint = torch.load(llama_ckpt_path, map_location="cpu")
     with open(Path(llama_weight_path) / "params.json", "r") as f:
         params = json.loads(f.read())