From ae1e80bce8cc7710266c5150d479bf56debf76f4 Mon Sep 17 00:00:00 2001
From: David Braun <2096055+DBraun@users.noreply.github.com>
Date: Sat, 30 Mar 2024 13:08:41 -0400
Subject: [PATCH] Update base.py

Restrict number of quantizers before creating DAC file
---
 dac/model/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dac/model/base.py b/dac/model/base.py
index 546b3cb..34510fa 100644
--- a/dac/model/base.py
+++ b/dac/model/base.py
@@ -215,6 +215,9 @@ def compress(
 
         codes = torch.cat(codes, dim=-1)
 
+        if n_quantizers is not None:
+            codes = codes[:, :n_quantizers, :]
+        
         dac_file = DACFile(
             codes=codes,
             chunk_length=chunk_length,
@@ -226,9 +229,6 @@ def compress(
             dac_version=SUPPORTED_VERSIONS[-1],
         )
 
-        if n_quantizers is not None:
-            codes = codes[:, :n_quantizers, :]
-
         self.padding = original_padding
         return dac_file