From afb12c397da6a593e925cd20d2c91a30c1cc14b8 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Fri, 9 Aug 2024 17:55:47 +0200
Subject: [PATCH] use make_linear_layer

---
 src/gluonts/torch/model/bim/module.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/gluonts/torch/model/bim/module.py b/src/gluonts/torch/model/bim/module.py
index 8ffde49640..c16cfb2ab1 100644
--- 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
         )