fix: move only VAE+vocoder to GPU during eval to prevent device mismatch
The previous check (next(feature_utils_orig.parameters()).device) only inspected the first parameter (from CLIP), missing CPU-stranded vocoder weights when the module was in a mixed-device state. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -124,15 +124,14 @@ class SelvaVaeRoundtrip:
|
||||
flush=True)
|
||||
|
||||
# Decode using model's feature_utils — same path as the sampler
|
||||
orig_device = next(feature_utils.parameters()).device
|
||||
if orig_device != device:
|
||||
feature_utils.to(device)
|
||||
tod = feature_utils.tod
|
||||
tod_orig_device = next(tod.parameters()).device
|
||||
tod.to(device)
|
||||
try:
|
||||
spec = feature_utils.decode(latent_unnorm)
|
||||
out = feature_utils.vocode(spec)
|
||||
finally:
|
||||
if orig_device != device:
|
||||
feature_utils.to(orig_device)
|
||||
tod.to(tod_orig_device)
|
||||
|
||||
out = out.float().cpu()
|
||||
if out.dim() == 1:
|
||||
|
||||
Reference in New Issue
Block a user