added avg_pool_k options to BWENet

xiph · Jun 25, 2024 · d9ae035 · d9ae035
1 parent 7a8a757
commit d9ae035
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/dnn/torch/osce/models/bwe_net.py b/dnn/torch/osce/models/bwe_net.py
@@ -75,7 +75,9 @@ def __init__(self,
  kernel_size32=15,
  kernel_size48=15,
  conv_gain_limits_db=[-12, 12],
- activation="AdaShape"
+ activation="AdaShape",
+ avg_pool_k32 = 8,
+ avg_pool_k48=12
  ):
 
  super().__init__()
@@ -99,8 +101,8 @@ def __init__(self,
 
  # non-linear transforms
  if activation == "AdaShape":
- self.tdshape1 = TDShaper(cond_dim, frame_size=self.frame_size32, avg_pool_k=8)
- self.tdshape2 = TDShaper(cond_dim, frame_size=self.frame_size48, avg_pool_k=12)
+ self.tdshape1 = TDShaper(cond_dim, frame_size=self.frame_size32, avg_pool_k=avg_pool_k32)
+ self.tdshape2 = TDShaper(cond_dim, frame_size=self.frame_size48, avg_pool_k=avg_pool_k48)
  self.act1 = self.tdshape1
  self.act2 = self.tdshape2
  elif activation == "ReLU":