Skip to content

Commit d7e1e71

Browse files
committed
fix: the exeption when topk is 1
1 parent ffa90e0 commit d7e1e71

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

inference.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ def main():
114114
_logger.info('Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
115115
batch_idx, len(loader), batch_time=batch_time))
116116

117-
topk_ids = np.concatenate(topk_ids, axis=0).squeeze()
117+
topk_ids = np.concatenate(topk_ids, axis=0)
118118

119119
with open(os.path.join(args.output_dir, './topk_ids.csv'), 'w') as out_file:
120120
filenames = loader.dataset.filenames(basename=True)

0 commit comments

Comments
 (0)