Update layers.py to ensure grad_output is contiguous (NVIDIA#1601)

Depending on how the forward pass is implemented, grad_output in this function will not be contiguous, so it crashes when calling .view(). By adding this extra line, we ensure it's contiguous. And it's a no-op when it already is contiguous, so performance shouldn't sufffer.

Author: markelsanz14

apex/transformer/tensor_parallel/layers.py

@@ -396,6 +396,7 @@ def backward(ctx, grad_output):
             return grad_input, None, None, None, None, None, None
 
         # Convert the tensor shapes to 2D for execution compatibility
+        grad_output = grad_output.contiguous()
         grad_output = grad_output.view(
             grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2]
         )