address #7

lucidrains · Apr 24, 2024 · a0d8a4e · a0d8a4e
1 parent e36ce95
commit a0d8a4e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'soft-moe-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.7',
+  version = '0.1.8',
   license='MIT',
   description = 'Soft MoE - Pytorch',
   author = 'Phil Wang',

diff --git a/soft_moe_pytorch/soft_moe.py b/soft_moe_pytorch/soft_moe.py
@@ -289,7 +289,10 @@ def __init__(
         super().__init__()
         assert exists(seq_len) ^ exists(num_slots), 'either seq_len, or num_slots must be passed into SoftMoE'
 
-        num_slots = default(num_slots, seq_len // num_experts)
+        if exists(seq_len):
+            num_slots = default(num_slots, seq_len // num_experts)
+        elif exists(num_slots):
+            seq_len = num_slots * num_experts
 
         norm_klass = LayerNorm if use_layernorm else RMSNorm
         self.norm = norm_klass(dim)