remove reshape in linear backward

alanwaketan · alanwaketan · commit 84b866173af7 · 2023-08-11T21:19:17.000Z
diff --git a/torch_xla/distributed/fsdp/utils.py b/torch_xla/distributed/fsdp/utils.py
@@ -120,7 +120,7 @@ def backward(ctx, grad_output):
 
 
 def _xla_patched_nn_linear_forward(m, input):
-  return XLAPatchedLinear.apply(input, m.weight, m.bias)
+  return XLAPatchedMatmul.apply(input, m.weight.t())
 
 
 def apply_xla_patch_to_nn_linear(module):
