From 164d3e4b36cb4017fb1c296bc0030dd1d38a7eb4 Mon Sep 17 00:00:00 2001 From: randaller Date: Sun, 5 Mar 2023 19:07:47 +0300 Subject: [PATCH] Update example-cpu.py --- example-cpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/example-cpu.py b/example-cpu.py index 7311308..3d915b9 100644 --- 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)