diff --git a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_mnist.py b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_mnist.py
index 0541c37fc71b08cde1f71fee72365d4535596a6f..09be10e6c8a7e9b676e434b410f702c3fe7bdb91 100644
--- a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_mnist.py
+++ b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_mnist.py
@@ -108,12 +108,9 @@ class MNIST(fluid.dygraph.Layer):
             loss = fluid.layers.cross_entropy(x, label)
             avg_loss = fluid.layers.mean(loss)
 
-        # TODO: Uncomment code after "return" statement can be transformed correctly.
-
-        #     return x, acc, avg_loss
-        # else:
-        #     return x
-        return x, acc, avg_loss
+            return x, acc, avg_loss
+        else:
+            return x
 
     def inference(self, inputs):
         x = self._simple_img_conv_pool_1(inputs)