From efb997d53138986615b802f10768bf5c319eae48 Mon Sep 17 00:00:00 2001 From: leyang Date: Wed, 8 May 2024 15:10:21 -0400 Subject: [PATCH] Fixed batch_size bug --- transfer_model/__main__.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/transfer_model/__main__.py b/transfer_model/__main__.py index 2ffc250..8d295df 100644 --- a/transfer_model/__main__.py +++ b/transfer_model/__main__.py @@ -87,11 +87,11 @@ def main() -> None: for ii, path in enumerate(paths): _, fname = osp.split(path) - - output_path = osp.join( - output_folder, f'{osp.splitext(fname)[0]}.pkl') - with open(output_path, 'wb') as f: - pickle.dump(var_dict, f) + if ii == 0: + output_path = osp.join( + output_folder, f'{osp.splitext(fname)[0]}.pkl') + with open(output_path, 'wb') as f: + pickle.dump(var_dict, f) output_path = osp.join( output_folder, f'{osp.splitext(fname)[0]}.obj')