From d7e1e7144a62f0368787c081eb999d27cb97d159 Mon Sep 17 00:00:00 2001 From: comar Date: Mon, 10 May 2021 16:52:40 +0900 Subject: [PATCH] fix: the exeption when topk is 1 --- inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference.py b/inference.py index 445dc4e5..5fcf1e60 100755 --- a/inference.py +++ b/inference.py @@ -114,7 +114,7 @@ def main(): _logger.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format( batch_idx, len(loader), batch_time=batch_time)) - topk_ids = np.concatenate(topk_ids, axis=0).squeeze() + topk_ids = np.concatenate(topk_ids, axis=0) with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file: filenames = loader.dataset.filenames(basename=True)