diff --git a/modules/estimators/nets.py b/modules/estimators/nets.py
index 5bd3debf..fb3dd128 100644
--- a/modules/estimators/nets.py
+++ b/modules/estimators/nets.py
@@ -61,10 +61,8 @@ def __init__(
         )
 
         # post process
-        # self.k_filter = nn.Parameter(torch.ones(num_speakers))
+        self.k_filter = nn.Parameter(torch.ones(num_speakers))
         self.p0 = (0.0 - self.vmin) / (self.vmax - self.vmin)
-        self.k_filter = nn.Embedding(num_speakers, 1)
-        nn.init.constant_(self.k_filter.weight, 0.0)
 
     def forward(self, x: torch.Tensor, spk_id: torch.Tensor=None) -> torch.Tensor:
         """
@@ -78,7 +76,7 @@ def forward(self, x: torch.Tensor, spk_id: torch.Tensor=None) -> torch.Tensor:
         x, _ = self.rnn(x)
         x = self.output_proj(x)
         if spk_id is not None:
-            k_filter = self.k_filter(spk_id).view(-1,1,1)
+            k_filter = self.k_filter[spk_id].view(-1,1,1)
             k_filter = 0.5 + 1.5 * torch.sigmoid(k_filter)
             k_filter = torch.clamp(k_filter, min=0.3, max=2.5)
             x = (x - self.p0) * k_filter + self.p0