fix(bigvgan-trainer): resolve device mismatch in _save_sample after offload

After the finally block, offload_to_cpu moves the vocoder to CPU while
ref_mel stays on GPU. Fix: detect vocoder's current device via
next(vocoder.parameters()).device and move ref_mel there before vocoding.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-09 01:35:07 +02:00
parent 9fdeb65182
commit 81ff0d46c9
+4 -1
View File
@@ -182,8 +182,11 @@ class SelvaBigvganTrainer:
def _save_sample(label):
"""Vocode the reference mel and save as .wav."""
try:
# Vocoder may have been offloaded to CPU after training — match its device.
voc_device = next(vocoder.parameters()).device
mel = ref_mel.to(voc_device)
with torch.no_grad():
wav = vocoder(ref_mel) # [1, 1, T] or [1, T]
wav = vocoder(mel) # [1, 1, T] or [1, T]
if wav.dim() == 2:
wav = wav.unsqueeze(1)
wav = wav.float().cpu().clamp(-1, 1)