From c99572a6d94a56a9f0dd2073ae405de23be28f7d Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Mon, 6 Jun 2022 11:44:19 +0200
Subject: [PATCH] fix args

---
 switch/module.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/switch/module.py b/switch/module.py
index 036d66a..ec42d30 100644
--- a/switch/module.py
+++ b/switch/module.py
@@ -170,10 +170,10 @@ class SwitchTransformerEncoderLayer(nn.Module):
         linear = nn.Linear(d_model, dim_feedforward, **factory_kwargs)
 
         self.linear1 = SwitchFeedForward(
-            capacity_factor,
-            drop_tokens,
-            is_scale_prob,
-            n_experts,
+            capacity_factor=capacity_factor,
+            drop_tokens=drop_tokens,
+            is_scale_prob=is_scale_prob,
+            n_experts=n_experts,
             expert=linear,
             d_model=d_model,
             dim_feedforward=dim_feedforward,