From 351d3ae0bd9909bcf1912e6f8882bdaf18e98f74 Mon Sep 17 00:00:00 2001
From: chenguowei01 <chenguowei01@baidu.com>
Date: Mon, 20 Jul 2020 20:17:04 +0800
Subject: [PATCH] update unet.py

---
 dygraph/models/unet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dygraph/models/unet.py b/dygraph/models/unet.py
index 970936d0..78d1a394 100644
--- a/dygraph/models/unet.py
+++ b/dygraph/models/unet.py
@@ -52,7 +52,7 @@ class UNet(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) / (
-- 
GitLab