Merge pull request databricks#43 from bcui19/removing_extra_size

Removing an extra size call

Author: tgale96

megablocks/layers/dmoe.py

@@ -316,7 +316,6 @@ def forward(self, x):
         # NOTE: If we're going to cast the activations to lower precision
         # do it before we permute the tokens to save bandwidth.
         x = common.cast_if_autocast_enabled(x)
-        sl, bs, hs = x.size()
 
         # Compute the expert scores and assignments.
         scores, expert_weights, top_experts = self.router(x)