cleanup

databricks · Aug 12, 2024 · 92b1607 · 92b1607
1 parent c14252d
commit 92b1607
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/megablocks/layers/mlp.py b/megablocks/layers/mlp.py
@@ -179,8 +179,7 @@ def create_dmoe_expert_weights(
         columns,
         init_method,
     )
-    weights = weights.view([-1, columns])
-    return weights
+    return weights.view([-1, columns])
 
 
 class MemoryOptimizedMLP(torch.autograd.Function):