[minor] use fp32 when generate on CPU

mit-han-lab · Nov 10, 2020 · c8d6d71 · c8d6d71
1 parent d5e13d0
commit c8d6d71
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/generate.py b/generate.py
@@ -31,6 +31,10 @@ def main(args):
     print(args)
 
     use_cuda = torch.cuda.is_available() and not args.cpu
+
+    # when running on CPU, use fp32 as default
+    if not use_cuda:
+        args.fp16 = False
 
     # Load dataset splits
     task = tasks.setup_task(args)