Spaces:

uto1125
/

Tai-Ya-test

Runtime error

uto1125 commited on Sep 21, 2024

Commit

cb462fd

verified ·

1 Parent(s): 271093a

Update tools/llama/generate.py

Files changed (1) hide show

tools/llama/generate.py CHANGED Viewed

@@ -341,7 +341,7 @@ def encode_tokens(
 def load_model(checkpoint_path, device, precision, compile=False):
     model: Union[NaiveTransformer, DualARTransformer] = BaseTransformer.from_pretrained(
-        checkpoint_path, load_weights=True
     )
     model = model.to(device=device, dtype=precision)

 def load_model(checkpoint_path, device, precision, compile=False):
     model: Union[NaiveTransformer, DualARTransformer] = BaseTransformer.from_pretrained(
+        checkpoint_path, load_weights=True, weights_only=True
     )
     model = model.to(device=device, dtype=precision)