Update example-cpu.py

3 years ago · 164d3e4b36
parent 7266efc685
commit 164d3e4b36
1 changed files with 1 additions and 0 deletions
--- a/example-cpu.py
+++ b/example-cpu.py
@ -53,6 +53,7 @@ def main(
        max_batch_size: int = 32,
 ):
    # torch.manual_seed(1)
+    # torch.set_default_dtype(torch.bfloat16)

    generator = load(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size)