properly specify hidden segments

jloveric · Dec 14, 2023 · e0fe41e · e0fe41e
1 parent a7af9a1
commit e0fe41e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/config/net/transformer.yaml b/config/net/transformer.yaml
@@ -19,7 +19,14 @@ inner: 10
 #layers: [[1, ${inner}, 128], [${inner}, ${inner}, 10], [${inner}, ${inner}, 10]]
 layers:
   [
-    { "input": 1, "hidden": 8, "output": 8, "layers": 1, "segments": 128 },
+    {
+      "input": 1,
+      "hidden": 8,
+      "output": 8,
+      "layers": 1,
+      "segments": 4,
+      "input_segments": 128,
+    },
     { "input": 8, "output": 8, "segments": 4 },
     { "input": 8, "output": 8, "segments": 4 },
     { "input": 8, "output": 8, "segments": 4 },

diff --git a/language_interpolation/networks.py b/language_interpolation/networks.py
@@ -346,6 +346,7 @@ def __init__(
         hidden_width = layers[0]["hidden"]
         embedding_layers = layers[0]["layers"]
         segments = layers[0]["segments"]
+        input_segments = layers[0]["input_segments"]
 
         mlp_normalization = None
         if normalization:
@@ -355,7 +356,7 @@ def __init__(
             layer_type=layer_type,
             n=n,
             in_width=input_width,
-            in_segments=segments,
+            in_segments=input_segments,
             out_segments=segments,
             hidden_segments=segments,
             hidden_layers=embedding_layers,