use make_linear_layer

awslabs · Aug 9, 2024 · afb12c3 · afb12c3
1 parent efa4f31
commit afb12c3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/gluonts/torch/model/bim/module.py b/src/gluonts/torch/model/bim/module.py
@@ -81,7 +81,9 @@ def __init__(
             self.scaler = NOPScaler(keepdim=True)
 
         # concat loc and scale to the context window
-        self.linear_backbone = nn.Linear(context_length + 2, self.mem_dim)
+        self.linear_backbone = make_linear_layer(
+            context_length + 2, self.mem_dim
+        )
         self.end_conv = nn.Conv1d(
             self.mem_dim * 3, self.mem_dim, kernel_size=1
         )