ALways export in fp32 base dtype on CPU

Xilinx · Jul 17, 2023 · 5fe532f · 5fe532f
1 parent a4881df
commit 5fe532f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/brevitas_examples/llm/main.py b/src/brevitas_examples/llm/main.py
@@ -291,6 +291,8 @@ def main():
 
     if args.export_target:
         print(f"Export to {args.export_target}")
+        # Currently we always export on CPU with a float32 container to avoid float16 CPU errors
+        model = model.cpu().to(dtype=torch.float32)
         model_export(model, calibration_loader[0], args)