Update README.md

README.md

@@ -34,7 +34,6 @@ import torch
 import IPython.display as ipd
 
 def redistribute_codes(row):
-    row = row[row > 128266]
     row_length = row.size(0)
     new_length = (row_length // 7) * 7
     trimmed_row = row[:new_length]
@@ -53,9 +52,14 @@ def redistribute_codes(row):
     
     with torch.no_grad():
 
-        codes = [torch.concat(layer_1)[None], 
-            torch.concat(layer_2)[None], 
-            torch.concat(layer_3)[None]]
+        codes = [torch.concat(layer_1), 
+            torch.concat(layer_2), 
+            torch.concat(layer_3)]
+        
+        for i in range(len(codes)):
+            codes[i][codes[i] < 0] = 0
+            codes[i] = codes[i][None]
+        
         audio_hat = snac_model.decode(codes)
         return audio_hat.cpu()[0, 0]