diff --git a/dygraph/models/hrnet.py b/dygraph/models/hrnet.py
index d5270a1089c46d390a05ca5c53df2d8e044a900c..2dcf2ddad0aad6e83356fe884ad213660b6816f3 100644
--- a/dygraph/models/hrnet.py
+++ b/dygraph/models/hrnet.py
@@ -179,6 +179,8 @@ class HRNet(fluid.dygraph.Layer):
             return pred, score_map
 
     def _get_loss(self, logit, label):
+        logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
+        label = fluid.layers.transpose(label, [0, 2, 3, 1])
         mask = label != self.ignore_index
         mask = fluid.layers.cast(mask, 'float32')
         loss, probs = fluid.layers.softmax_with_cross_entropy(
@@ -186,7 +188,7 @@ class HRNet(fluid.dygraph.Layer):
             label,
             ignore_index=self.ignore_index,
             return_softmax=True,
-            axis=1)
+            axis=-1)
 
         loss = loss * mask
         avg_loss = fluid.layers.mean(loss) / (
diff --git a/dygraph/models/unet.py b/dygraph/models/unet.py
index 9bb92302f5a65eac6f536f0308f01ea878361aa7..970936d05e0fb67a58973879ffceef5f02816495 100644
--- a/dygraph/models/unet.py
+++ b/dygraph/models/unet.py
@@ -43,6 +43,8 @@ class UNet(fluid.dygraph.Layer):
             return pred, score_map
 
     def _get_loss(self, logit, label):
+        logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
+        label = fluid.layers.transpose(label, [0, 2, 3, 1])
         mask = label != self.ignore_index
         mask = fluid.layers.cast(mask, 'float32')
         loss, probs = fluid.layers.softmax_with_cross_entropy(