From 400858834394b830cce314971f6826a1b89dab81 Mon Sep 17 00:00:00 2001
From: gaotingquan <gaotingquan@baidu.com>
Date: Tue, 7 Mar 2023 11:40:58 +0000
Subject: [PATCH] add amp decorator and parse_batch decorator

---
 ppcls/arch/backbone/legendary_models/mobilenet_v1.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
index 3e612068..499fad7f 100644
--- a/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
+++ b/ppcls/arch/backbone/legendary_models/mobilenet_v1.py
@@ -22,6 +22,8 @@ from paddle.nn import Conv2D, BatchNorm, Linear, ReLU, Flatten
 from paddle.nn import AdaptiveAvgPool2D
 from paddle.nn.initializer import KaimingNormal
 
+from ..base import clas_forward_decorator
+from ....utils.amp import AMP_forward_decorator
 from ..base.theseus_layer import TheseusLayer
 from ....utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
 
@@ -162,6 +164,8 @@ class MobileNet(TheseusLayer):
             return_patterns=return_patterns,
             return_stages=return_stages)
 
+    @AMP_forward_decorator
+    @clas_forward_decorator
     def forward(self, x):
         x = self.conv(x)
         x = self.blocks(x)
-- 
GitLab