Update modeling_super_linear.py

modeling_super_linear.py

@@ -347,7 +347,7 @@ class SparseNoisyMoE(nn.Module):
             x_0 = x
             
         self.gate_outputs = self.gating_network(x_0)
-        print(self.gate_outputs.shape)
+        #print(self.gate_outputs.shape)
 
         if not self.training:
             self.gate_outputs = self.gate_outputs / self.moe_temp
@@ -519,7 +519,8 @@ class superLinear(nn.Module):
             x = x_enc.permute(0, 2, 1)
             B, V, L = x.shape
         else:
-            x    = x_enc
+            x     = x_enc
+            x_enc =  x_enc.unsqueeze(-1)
             B, L = x.shape
             V    = 1