src/biencoder_predict_qa.py [196:203]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    heads = {}
    for name in ['start', 'end']:
        head = MockClassificationHead.load_from_file(
                os.path.join(args.load_dir, f'model_qa_head_{name}.pt'),
                do_token_level=True)
        head.to('cuda')
        heads[name] = head
    print('Finished loading model.')
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



src/run_sentiment.py [129:136]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        heads = {}
        for name in ['start', 'end']:
            head = MockClassificationHead.load_from_file(
                    os.path.join(args.load_dir, f'model_qa_head_{name}.pt'),
                    do_token_level=True)
            head.to('cuda')
            heads[name] = head
    print('Finished loading model.')
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -