diff --git a/metadata/save_meta.py b/metadata/save_meta.py
index 8a83b9e..7f6ed5d 100644
--- a/metadata/save_meta.py
+++ b/metadata/save_meta.py
@@ -6,9 +6,10 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--data_path', type=str)
     parser.add_argument('--json_path', type=str)
+    parser.add_argument('--extension', type=str, default='pt')
     args = parser.parse_args()
 
-    fpath_list = sorted([os.path.join(args.data_path, fname) for fname in os.listdir(args.data_path) if fname.endswith('.pt')])
+    fpath_list = sorted([os.path.join(args.data_path, fname[:-len(args.extension)-1]) for fname in os.listdir(args.data_path) if fname.endswith('.' + args.extension)])
     os.makedirs(args.json_path, exist_ok=True)
     json.dump(fpath_list, open(args.json_path, 'w'))