| 149 | |
| 150 | class ConvTBC(nn.Module): |
| 151 | def __init__(self, in_channels, out_channels, kernel_size, padding=0): |
| 152 | super(ConvTBC, self).__init__() |
| 153 | self.in_channels = in_channels |
| 154 | self.out_channels = out_channels |
| 155 | self.kernel_size = kernel_size |
| 156 | self.padding = padding |
| 157 | |
| 158 | self.weight = torch.nn.Parameter(torch.Tensor( |
| 159 | self.kernel_size, in_channels, out_channels)) |
| 160 | self.bias = torch.nn.Parameter(torch.Tensor(out_channels)) |
| 161 | |
| 162 | def forward(self, input): |
| 163 | return torch.conv_tbc(input.contiguous(), self.weight, self.bias, self.padding) |