def METHOD_NAME(gpu_idx: int):
    dataset = load_dataset('json', data_files=kd_gen_score_in_path)['train']
    if args.dry_run:
        dataset = dataset.select(range(100))
    dataset = dataset.shard(num_shards=torch.cuda.device_count(),
                            index=gpu_idx,
                            contiguous=True)
    qid_pids = []
    for ex in tqdm.tqdm(dataset, desc='get qid-pid pairs', mininterval=3):
        for pos_doc_id in ex['positives']['doc_id']:
            qid_pids.append((ex['query_id'], pos_doc_id))
        for neg_doc_id in ex['negatives']['doc_id'][:args.kd_gen_score_n_neg]:
            qid_pids.append((ex['query_id'], neg_doc_id))
    dataset = Dataset.from_dict({'query_id': [t[0] for t in qid_pids],
                                 'doc_id': [t[1] for t in qid_pids]})
    query_ids, doc_ids = dataset['query_id'], dataset['doc_id']
    logger.info('GPU {} needs to process {} examples'.format(gpu_idx, len(dataset)))
    torch.cuda.set_device(gpu_idx)
    tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(args.model_name_or_path)
    model: RerankerForInference = RerankerForInference.from_pretrained(args.model_name_or_path)
    model.eval()
    model.cuda()
    corpus: Dataset = load_corpus(path=os.path.join(args.data_dir, 'passages.jsonl.gz'))
    queries = load_queries(path='{}/{}_queries.tsv'.format(args.data_dir, args.kd_gen_score_split),
                           task_type=args.task_type)
    dataset.set_transform(partial(_kd_gen_score_transform_func, tokenizer, corpus, queries))
    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args.fp16 else None)
    data_loader = DataLoader(
        dataset,
        batch_size=args.kd_gen_score_batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=args.dataloader_num_workers,
        collate_fn=data_collator,
        pin_memory=True)
    scores = []
    for batch_dict in tqdm.tqdm(data_loader, desc='generate teacher score', mininterval=5):
        batch_dict = move_to_cuda(batch_dict)
        with torch.cuda.amp.autocast() if args.fp16 else nullcontext():
            outputs: SequenceClassifierOutput = model(batch_dict)
        scores.append(outputs.logits.squeeze(dim=-1).cpu())
        assert len(scores[-1].shape) == 1
    all_scores = torch.cat(scores, dim=-1)
    assert all_scores.shape[0] == len(dataset), '{} != {}'
    all_scores = all_scores.tolist()
    with open(_get_shard_path(gpu_idx), 'w', encoding='utf-8') as writer:
        for idx in range(len(query_ids)):
            writer.write('{}\t{}\t{}\n'.format(query_ids[idx], doc_ids[idx], round(all_scores[idx], 5)))
    logger.info('Done computing teacher score for worker {}'.format(gpu_idx))