Upload folder using huggingface_hub

plpq.py

@@ -83,9 +83,9 @@ class PLPQ(PreTrainedModel):
         """
         Quantize the input tensor
         Parameters:
-            x (Image or torch.Tensor): The input tensor. Size b, c, h, w
+            x (torch.Tensor): The input tensor of shape (b, c, h, w)
         Returns:
-            torch.Tensor: The indices tensor. Size b, h, w
+            torch.Tensor: The indices tensor of shape (b, t, n_quantizers)
         """
         z = self.encoder(x).permute(0, 2, 3, 1).contiguous()
         b, h, w, c = z.shape
@@ -97,10 +97,11 @@ class PLPQ(PreTrainedModel):
     @torch.no_grad()
     def decode(self, indices: torch.Tensor) -> torch.Tensor:
         """
-        Parameters: 
-            indices: torch.Tensor of shape (b, t, n_freq_bins)
-        Returns: 
-            emb: torch.Tensor of shape (b, t, n_embd)
+        Decode a tensor, inverse of self.quantize
+        Parameters:
+            indices (torch.Tensor): The input codes of shape (b, t, n_quantizers)
+        Returns:
+            torch.Tensor: The decoded tensor of shape (b, c, h, w)
         """
 
         ncodes = indices.shape[-1]