From 0f2dda78978b65f5cc02f0b2632039fd0a5f27d8 Mon Sep 17 00:00:00 2001
From: zk1998 <zhoukun_1998@qq.com>
Date: Mon, 30 Sep 2024 16:58:01 +0800
Subject: [PATCH] [tracer] skip_unsupport_torchvision_op

---
 tinynn/graph/configs/gen_modules_yml.py        | 8 +++++++-
 tinynn/graph/configs/torch_module_override.yml | 3 ---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/tinynn/graph/configs/gen_modules_yml.py b/tinynn/graph/configs/gen_modules_yml.py
index a00453e3..ddc0fbe4 100644
--- a/tinynn/graph/configs/gen_modules_yml.py
+++ b/tinynn/graph/configs/gen_modules_yml.py
@@ -49,7 +49,13 @@
         print(k, c, issubclass(c, torch.nn.Module))
         print(c.__module__)
         # Skip container modules
-        if c.__name__ in ('FeaturePyramidNetwork', 'SqueezeExcitation'):
+        if c.__name__ in (
+            'FeaturePyramidNetwork',
+            'SqueezeExcitation',
+            'Conv2dNormActivation',
+            'Conv3dNormActivation',
+            'MLP',
+        ):
             continue
         final_dict.setdefault('torchvision.ops', [])
         final_dict['torchvision.ops'].append(k)
diff --git a/tinynn/graph/configs/torch_module_override.yml b/tinynn/graph/configs/torch_module_override.yml
index 2f55ac5a..6513ac8c 100644
--- a/tinynn/graph/configs/torch_module_override.yml
+++ b/tinynn/graph/configs/torch_module_override.yml
@@ -159,10 +159,7 @@ torchvision.ops:
 - DeformConv2d
 - DropBlock2d
 - DropBlock3d
-- Conv2dNormActivation
-- Conv3dNormActivation
 - FrozenBatchNorm2d
-- MLP
 - Permute
 - MultiScaleRoIAlign
 - PSRoIAlign