Трансформаторы NLP - PyTorch RuntimeError: shape [5, -1, 32783] недопустим для ввода размером 655360

Question

Я пытаюсь загрузить модель SuperTransformer один раз и выполнить прогнозирование набора данных несколько раз в цикле, при этом каждый раз меняя конфигурацию выборочной модели.

Я использую файл generate.py в github.com/mit-han-lab/hardware-aware-transformers для выполнения переводов в наборе данных WMT'14En-De. Часть кода, изменяющая конфигурацию выборочной модели: model.set_sample_config(config) в строке 67 файла generate.py.

По сути, я пытаюсь повторить вывод для набора данных с разными конфигурациями субтрансформатора - при этом загружая набор данных и модель только один раз (сокращая как можно больше накладных расходов на второй итерации).

В настоящее время я пытаюсь реализовать измененный код ниже, с циклом while после этапа ансамбля #Load и изменением переменной config после каждой итерации.

       # Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, _model_args = checkpoint_utils.load_model_ensemble(
    args.path.split(':'),
    arg_overrides=eval(args.model_overrides),
    task=task,
)

torch.manual_seed(args.seed)

#loop twice, change 'config' variable on second iteration to have different SubTransformer configuration
while i in range (0,2):

    # Optimize ensemble for generation
    for model in models:
        if use_cuda:
            model.cuda()

        config = utils.get_subtransformer_config(args)

        model.set_sample_config(config)
        model.make_generation_fast_(
            beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
            need_attn=args.print_alignment,
        )

Однако я получаю следующую ошибку: форма недопустима для ввода размера x

Проблема, похоже, связана с размером переменной lprobs в sequence_generator.py, где размер всегда неверен на второй итерации. Эта ошибка возникает даже тогда, когда код снова зацикливается с той же конфигурацией субтрансформатора.

Я новичок в библиотеках PyTorch и в машинном обучении в целом, поэтому я буду благодарен за любую помощь, которую я могу получить с этим! Спасибо.