From 92b1607dcd822766a19c01196f5ebfcc29870173 Mon Sep 17 00:00:00 2001
From: Eitan Turok <eitan.turok@databricks.com>
Date: Mon, 12 Aug 2024 18:11:35 +0000
Subject: [PATCH] cleanup

---
 megablocks/layers/mlp.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/megablocks/layers/mlp.py b/megablocks/layers/mlp.py
index 5f73dd3..f7cb782 100644
--- 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):