Preserve the existing branch matrix axis

token_operations.py

@@ -72,7 +72,8 @@ def install(model,*,linears,moe):
             if module.weight.dtype!=torch.bfloat16 or (is_output and value.ndim==2):
                 return original(value,*args,**kwargs)
             shape=value.shape
-            result=batched_linear(value.reshape(1,-1,shape[-1]),module.weight,module.bias)
+            matrices=value if value.ndim==3 else value.reshape(1,-1,shape[-1])
+            result=batched_linear(matrices,module.weight,module.bias)
             return result.reshape(*shape[:-1],module.weight.shape[0])
         return forward
     if linears: