Fix DDP parameter and gradient layouts

backbones.py

@@ -151,6 +151,32 @@ class AgeSinusoidalEncoding(nn.Module):
         return output
 
 
+class DepthwiseConv2d(nn.Module):
+    """Depthwise convolution without a singleton weight dimension."""
+
+    def __init__(self, channels: int, kernel_size: int):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.padding = kernel_size // 2
+        self.weight = nn.Parameter(
+            torch.empty(channels, kernel_size, kernel_size)
+        )
+        self.bias = nn.Parameter(torch.empty(channels))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        bound = 1 / math.sqrt(kernel_size * kernel_size)
+        nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.conv2d(
+            x,
+            self.weight.unsqueeze(1),
+            self.bias,
+            padding=self.padding,
+            groups=self.channels,
+        )
+
+
 class LiteTimesBackbone2d(nn.Module):
     """Cheap local feature extractor for a TimesNet period image."""
 
@@ -162,10 +188,7 @@ class LiteTimesBackbone2d(nn.Module):
         if expansion <= 0:
             raise ValueError("expansion must be > 0")
         hidden_dim = max(dim, int(round(dim * expansion)))
-        self.dwconv = nn.Conv2d(
-            dim, dim, kernel_size=kernel_size,
-            padding=kernel_size // 2, groups=dim,
-        )
+        self.dwconv = DepthwiseConv2d(dim, kernel_size)
         self.norm = nn.GroupNorm(1, dim)
         self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
         self.act = nn.GELU()