Gumbel fixes

custom_interface.py

@@ -1,3 +1,4 @@
+from typing import Mapping
 import torch
 import math
 from speechbrain.inference.interfaces import Pretrained
@@ -80,7 +81,6 @@ class Discrete_EmbeddingLayer(torch.nn.Module):
         self.layers = layers
         self.available_layers = available_layers
         self.offsets = self.build_offsets()
-        self.layer_embs = self.compute_layer_embs()
         self.chunk_size = chunk_size
 
     def init_embedding(self, weights):
@@ -95,7 +95,10 @@ class Discrete_EmbeddingLayer(torch.nn.Module):
         )
         if self.layers:
             selected_layers = set(self.layers)
-            indexes = [idx for idx, layer in enumerate(self.layers) if layer in selected_layers]
+            indexes = [
+                idx for idx, layer in enumerate(self.available_layers)
+                if layer in selected_layers
+            ]
             offsets = offsets[indexes]
         return offsets
 
@@ -190,6 +193,11 @@ class Discrete_EmbeddingLayer(torch.nn.Module):
         )
         return emb
 
+    def load_state_dict(self, state_dict, strict=True):
+        result = super().load_state_dict(state_dict, strict)
+        self.layer_embs = self.compute_layer_embs()
+        return result
+
 
 class DiscreteSpkEmb(Pretrained):
     """A ready-to-use class for utterance-level classification (e.g, speaker-id,