From 600d00db7a814bd44e8d6eabf734b1fd7a11e57b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20C=2E=20Santos?= Date: Mon, 4 Sep 2023 13:29:31 +0100 Subject: [PATCH] Added a flag to export.py to allow exporting the best checkpoint and not just the last checkpoint --- scripts/export.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/export.py b/scripts/export.py index 40a4f0d..785c5d3 100644 --- a/scripts/export.py +++ b/scripts/export.py @@ -28,6 +28,10 @@ default=None, help='Path to the run to export', required=True) +flags.DEFINE_string('mode', + default='last', + help='Choose between exporting the last checkpoint or the best. Default is "last"', + required=False) flags.DEFINE_bool('streaming', default=False, help='Enable the model streaming mode') @@ -338,7 +342,7 @@ def main(argv): logging.info("building rave") gin.parse_config_file(os.path.join(FLAGS.run, "config.gin")) - checkpoint = rave.core.search_for_run(FLAGS.run) + checkpoint = rave.core.search_for_run(FLAGS.run, FLAGS.mode) pretrained = rave.RAVE() if checkpoint is not None: