diff --git a/autodist/kernel/synchronization/compressor.py b/autodist/kernel/synchronization/compressor.py
index 868b542..917f553 100644
--- a/autodist/kernel/synchronization/compressor.py
+++ b/autodist/kernel/synchronization/compressor.py
@@ -300,5 +300,5 @@ def _modified_gram_schmidt(matrix):
 
             rest = matrix[:, (i + 1):]
             rest -= math_ops.reduce_sum_v1(v * rest, axis=0, keepdims=True) * v
-            matrix = array_ops.concat([matrix[:, :i], v, rest], axis=1)
+            matrix = array_ops.concat([matrix[:, :i], v, rest], 1)
         return matrix