From 8ff01e850d667187d3e2f929693ac2d45e6b628c Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Sat, 28 May 2022 22:41:48 -0300 Subject: [PATCH] Fix gen_wts_yolor --- utils/gen_wts_yolor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/utils/gen_wts_yolor.py b/utils/gen_wts_yolor.py index 7f19da9..95def7a 100644 --- a/utils/gen_wts_yolor.py +++ b/utils/gen_wts_yolor.py @@ -20,7 +20,7 @@ def parse_args(): pt_file, cfg_file = parse_args() -wts_file = cfg_file.split(".cfg")[0] + ".wts" +wts_file = "%s.wts" % cfg_file.rsplit("/")[1].split(".cfg")[0] device = select_device("cpu") model = Darknet(cfg_file).to(device) @@ -41,3 +41,5 @@ with open(wts_file, "w") as f: conv_count += 1 f.write("{}\n".format(conv_count)) f.write(wts_write) + +os.system("cp %s ./" % cfg_file)