Spaces:

chansung
/

LLaMA-7B

Runtime error

chansung commited on Mar 3, 2023

Commit

33cc221

•

1 Parent(s): 32d38c1

Update gen.py

Files changed (1) hide show

gen.py CHANGED Viewed

@@ -71,7 +71,7 @@ def get_pretrained_models(
     tokenizer = Tokenizer(model_path=f"{tokenizer_weight_path}/tokenizer.model")
     model_args.vocab_size = tokenizer.n_words
     torch.set_default_tensor_type(torch.cuda.HalfTensor)
-    model = Transformer(model_args)
     torch.set_default_tensor_type(torch.FloatTensor)
     model.load_state_dict(checkpoint, strict=False)

     tokenizer = Tokenizer(model_path=f"{tokenizer_weight_path}/tokenizer.model")
     model_args.vocab_size = tokenizer.n_words
     torch.set_default_tensor_type(torch.cuda.HalfTensor)
+    model = Transformer(model_args).cuda().half()
     torch.set_default_tensor_type(torch.FloatTensor)
     model.load_state_dict(checkpoint, strict=False)