fix conv1d->conv2d in STNorm (#117)

GestaltCogTeam · Jan 4, 2024 · f001944 · f001944
1 parent b611d86
commit f001944
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/baselines/STNorm/arch/stnorm_arch.py b/baselines/STNorm/arch/stnorm_arch.py
@@ -104,16 +104,16 @@ def __init__(self, num_nodes, tnorm_bool, snorm_bool, in_dim, out_dim, channels,
                 self.filter_convs.append(nn.Conv2d(
                     in_channels=num * channels, out_channels=channels, kernel_size=(1, kernel_size), dilation=new_dilation))
 
-                self.gate_convs.append(nn.Conv1d(
+                self.gate_convs.append(nn.Conv2d(
                     in_channels=num * channels, out_channels=channels, kernel_size=(1, kernel_size), dilation=new_dilation))
 
                 # 1x1 convolution for residual connection
                 self.residual_convs.append(
-                    nn.Conv1d(in_channels=channels, out_channels=channels, kernel_size=(1, 1)))
+                    nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=(1, 1)))
 
                 # 1x1 convolution for skip connection
                 self.skip_convs.append(
-                    nn.Conv1d(in_channels=channels, out_channels=channels, kernel_size=(1, 1)))
+                    nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=(1, 1)))
                 new_dilation *= 2
                 receptive_field += additional_scope
                 additional_scope *= 2