debug

revq/models/backbone/dcae.py

@@ -240,6 +240,7 @@ class DCAEEncoder(nn.Module):
 class DCAEDecoder(nn.Module):
     def __init__(
         self,
+        TYPE: str,
         in_channels: int,
         latent_channels: int,
         attention_head_dim: int = 32,

revq/models/revq_quantizer.py

@@ -15,7 +15,8 @@ def find_nearest(x, y):
     return min_dist, indices
 
 class Quantizer(nn.Module):
-    def __init__(self, code_dim: int = 128, num_code: int = 1024,
+    def __init__(self, TYPE: str = "vq", 
+                 code_dim: int = 4, num_code: int = 1024,
                  num_group: int = None, tokens_per_data: int = 4,
                  auto_reset: bool = True, reset_ratio: float = 0.2,
                  reset_noise: float = 0.1