diff --git a/src/crfrnn_layer.py b/src/crfrnn_layer.py
index 1170e6f..cf34845 100644
--- a/src/crfrnn_layer.py
+++ b/src/crfrnn_layer.py
@@ -88,7 +88,7 @@ def call(self, inputs):
         q_values = unaries
 
         for i in range(self.num_iterations):
-            softmax_out = tf.nn.softmax(q_values, dim=0)
+            softmax_out = tf.nn.softmax(q_values, axis=0)
 
             # Spatial filtering
             spatial_out = custom_module.high_dim_filter(softmax_out, rgb, bilateral=False,