address #135 again

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.14.25"
+version = "1.14.26"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch/residual_vq.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import random
 from math import ceil
 from functools import partial
 from itertools import zip_longest
+from typing import List
 
 import torch
-from torch import nn
+from torch import nn, Tensor
 import torch.nn.functional as F
 from vector_quantize_pytorch.vector_quantize_pytorch import VectorQuantize
 
@@ -122,7 +125,7 @@ def forward(
         self,
         x,
         mask = None,
-        indices = None,
+        indices: Tensor | List[Tensor] | None = None,
         return_all_codes = False,
         sample_codebook_temp = None,
         freeze_codebook = False,
@@ -140,6 +143,9 @@ def forward(
         all_losses = []
         all_indices = []
 
+        if isinstance(indices, list):
+            indices = torch.stack(indices)
+
         if return_loss:
             assert not torch.any(indices == -1), 'some of the residual vq indices were dropped out. please use indices derived when the module is in eval mode to derive cross entropy loss'
             ce_losses = []