Updated Code to run on Colab Notebook without Errors

README.md

@@ -153,6 +153,9 @@ def decode_snac_tokens(snac_tokens, snac_model):
     if not snac_tokens or len(snac_tokens) % 7 != 0:
         return None
 
+    # Get the device of the SNAC model. Fixed by Shresth to run on colab notebook :)
+    snac_device = next(snac_model.parameters()).device
+
     # De-interleave tokens into 3 hierarchical levels
     codes_lvl = [[] for _ in range(3)]
     llm_codebook_offsets = [AUDIO_CODE_BASE_OFFSET + i * 4096 for i in range(7)]
@@ -172,7 +175,7 @@ def decode_snac_tokens(snac_tokens, snac_model):
     # Convert to tensors for SNAC decoder
     hierarchical_codes = []
     for lvl_codes in codes_lvl:
-        tensor = torch.tensor(lvl_codes, dtype=torch.int32, device=snac_model.device).unsqueeze(0)
+        tensor = torch.tensor(lvl_codes, dtype=torch.int32, device=snac_device).unsqueeze(0)
         if torch.any((tensor < 0) | (tensor > 4095)):
             raise ValueError("Invalid SNAC token values")
         hierarchical_codes.append(tensor)