fix snac

utils/snac_utils.py

@@ -21,8 +21,8 @@ def layershift(input_id, layer, stride=4160, shift=152000):
     return input_id + shift + layer * stride
 
     
-def generate_audio_data(snac_tokens, snacmodel):
-    audio = reconstruct_tensors(snac_tokens)
+def generate_audio_data(snac_tokens, snacmodel, device=None):
+    audio = reconstruct_tensors(snac_tokens, device)
     with torch.inference_mode():
         audio_hat = snacmodel.decode(audio)
     audio_data = audio_hat.cpu().numpy().astype(np.float64) * 32768.0