add modules

nod-ai · Nov 22, 2024 · 40fbd1a · 40fbd1a
1 parent 0e100e9
commit 40fbd1a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/sharktank/layers/modulation.py b/sharktank/sharktank/layers/modulation.py
@@ -20,7 +20,7 @@ def __init__(self, theta: Theta, double: bool):
 
         self.is_double = double
         self.multiplier = 6 if double else 3
-        self.lin = LinearLayer(theta("lin"))
+        self.add_module("lin", LinearLayer(theta("lin")))
 
     def forward(self, vec: torch.Tensor) -> tuple[ModulationOut, ModulationOut | None]:
         silu_result = ops.elementwise(F.silu, vec)