From b55947595eb973cda449f9a4552272efa6a5a345 Mon Sep 17 00:00:00 2001
From: 0x45f <23097963+0x45f@users.noreply.github.com>
Date: Fri, 19 Nov 2021 16:48:07 +0800
Subject: [PATCH] set net.forward to original forward function in flops
 (#36852) (#37357)

set net.forward to original forward function in flops when net is a dy2stat model.
---
 python/paddle/hapi/dynamic_flops.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/paddle/hapi/dynamic_flops.py b/python/paddle/hapi/dynamic_flops.py
index 2c59ee67d4a..5c48501c07d 100644
--- a/python/paddle/hapi/dynamic_flops.py
+++ b/python/paddle/hapi/dynamic_flops.py
@@ -17,6 +17,7 @@ import warnings
 import paddle.nn as nn
 import numpy as np
 from .static_flops import static_flops, Table
+from paddle.fluid.dygraph.dygraph_to_static.program_translator import unwrap_decorators
 
 __all__ = []
 
@@ -100,6 +101,10 @@ def flops(net, input_size, custom_ops=None, print_detail=False):
             #Total Flops: 347560     Total Params: 61610
     """
     if isinstance(net, nn.Layer):
+        # If net is a dy2stat model, net.forward is StaticFunction instance,
+        # we set net.forward to original forward function.
+        _, net.forward = unwrap_decorators(net.forward)
+
         inputs = paddle.randn(input_size)
         return dynamic_flops(
             net,
-- 
GitLab