From 8a5e2043bbd2d0995037a3644c2e3b5386b26130 Mon Sep 17 00:00:00 2001 From: Yiqing-Zhou Date: Tue, 9 May 2023 00:37:52 +0800 Subject: [PATCH] [optimize] map_location='cpu' for load_from_checkpoint --- lit_export.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lit_export.py b/lit_export.py index 1f03924..07cac08 100644 --- a/lit_export.py +++ b/lit_export.py @@ -26,6 +26,8 @@ if __name__ == '__main__': checkpoint_file_path = next(lightning_logs_dir_path.glob("checkpoints/*.ckpt")) - lit_module = LitModule.load_from_checkpoint(checkpoint_file_path) + lit_module = LitModule.load_from_checkpoint( + checkpoint_file_path, map_location='cpu' + ) model: PreTrainedModel = lit_module.__core_module__ model.save_pretrained(exports_dir_path)