Update router.py (#158)

databricks · Oct 17, 2024 · 84286de · 84286de
1 parent 7b0337f
commit 84286de
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/megablocks/layers/router.py b/megablocks/layers/router.py
@@ -79,9 +79,8 @@ def __init__(self, args: Arguments):
         args.init_method(self.layer.weight)
 
     def jitter(self, x: torch.Tensor):
-        assert isinstance(self.args.moe_jitter_eps, float)
-        low = 1.0 - self.args.moe_jitter_eps
-        high = 1.0 + self.args.moe_jitter_eps
+        low: float = 1.0 - self.args.moe_jitter_eps
+        high: float = 1.0 + self.args.moe_jitter_eps
         noise = torch.rand(x.size(), dtype=x.dtype, device=x.device)
         return low + noise * (high - low)