Conformer specify conv_layers type

rwth-i6 · Apr 4, 2023 · 6182491 · 6182491
1 parent eed07b5
commit 6182491
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nn/conformer.py b/nn/conformer.py
@@ -121,7 +121,7 @@ def __init__(
         self.pool_sizes = pool_sizes
         self.activation = activation
 
-        self.conv_layers = nn.ModuleList()
+        self.conv_layers: nn.ModuleList[nn.Conv2d] = nn.ModuleList()
         if strides is None:
             strides = [1] * len(out_dims)
         assert len(out_dims) == len(filter_sizes) == len(strides) > 0