PhysiQuanty
/

Binary-LLM-POC

Text Generation

TinyTransformerLM

Model card Files Files and versions

PhysiQuanty commited on Feb 2

Commit

ca1b436

·

verified ·

1 Parent(s): 4571215

Update inference.py

Files changed (1) hide show

inference.py +1 -1

inference.py CHANGED Viewed

@@ -198,7 +198,7 @@ def main() -> None:
 	with torch.no_grad():
 		for _ in range(int(args.max_new_tokens)):
 			# full forward sur toute la séquence, sans cache
-			out = m(input_ids=tokens, use_cache=False)
 			logits = out.logits[:, -1, :]
 			full_seq = tokens[0].tolist()

 	with torch.no_grad():
 		for _ in range(int(args.max_new_tokens)):
 			# full forward sur toute la séquence, sans cache
+			out = m(input_ids=tokens, use_cache=True)
 			logits = out.logits[:, -1, :]
 			full_seq = tokens[0].tolist()