Commit
•
13b9094
1
Parent(s):
ecec5b7
Fix (minimal): Move to CPU before export
Browse files- minimal_script.py +1 -0
minimal_script.py
CHANGED
|
@@ -205,6 +205,7 @@ def main(args):
|
|
| 205 |
torch.save(pipe.unet.state_dict(), os.path.join(output_dir, args.checkpoint_name))
|
| 206 |
|
| 207 |
if args.export_target:
|
|
|
|
| 208 |
export_quant_params(pipe, output_dir)
|
| 209 |
|
| 210 |
|
|
|
|
| 205 |
torch.save(pipe.unet.state_dict(), os.path.join(output_dir, args.checkpoint_name))
|
| 206 |
|
| 207 |
if args.export_target:
|
| 208 |
+
pipe.unet.to('cpu').to(dtype)
|
| 209 |
export_quant_params(pipe, output_dir)
|
| 210 |
|
| 211 |
|