GPU inference GigaAM-RNNT in float16 dtype

[potassium-chloride]

Добрый день! В моей видеокарте мало памяти и модель умещается только если хранить её с понижением точности в float16, однако в таком случае не работает transcribe. Мои версии библиотек:

torch==2.3.1+cu121
torchaudio==2.3.1+cu121
nemo_toolkit @ git+https://github.com/NVIDIA/NeMo.git@1fa961ba03ab5f8c91b278640e29807079373372

Я попытался просто перед переносом на видеокарту вызвать .to(torch.float16) и обернуть model.transcribe в torch.autocast(dtype=torch.float16), но это вызвало ошибку:

torchaudio/transforms/_transforms.py", line 412, in forward
    mel_specgram = torch.matmul(specgram.transpose(-1, -2), self.fb).transpose(-1, -2)
RuntimeError: expected scalar type Half but found Float

Тогда я перегрузил метод forward класса FilterbankFeaturesTA и вставил между signals = self._apply_preemphasis(signals=signals) и features = self._extract_spectrograms(signals=signals) явное преобразование signals = signals.to(torch.float16) и отредакторовал yaml файл конфигурации, сделав n_fft равным 512, так как иначе не выполняется FFT на видеокарте, но получаю в качестве транскрипции множество сдвоенных вопросительных знаков. При этом я отдельно на CPU проверил, что изменение размера окна преобразования Фурье в препроцессоре не влияет на качество распознавания речи (по крайней мере на моих примерах).

При этом если не делать with torch.autocast(dtype=torch.float16), то код вываливается с ошибкой в энкодере:

.../nemo/collections/asr/models/rnnt_models.py", line 670, in forward
    encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
...
    return F.conv1d(input, weight, bias, self.stride,
RuntimeError: Input type (float) and bias type (c10::Half) should be the same

Существует ли способ выполнять вычисления в float16?