support flatten (#4628)

0a52ae6f · Ningxin Zheng · GitHub · 68ca6f21 · 0a52ae6f
Unverified Commit 0a52ae6f authored Mar 11, 2022 by Ningxin Zheng Committed by GitHub Mar 11, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

nni/compression/pytorch/speedup/compress_modules.py nni/compression/pytorch/speedup/compress_modules.py +2 -1

No files found.
--- a/nni/compression/pytorch/speedup/compress_modules.py
+++ b/nni/compression/pytorch/speedup/compress_modules.py
@@ -41,7 +41,8 @@ replace_module = {
    'Dropout3d': lambda module, masks: no_replace(module, masks),
    'Upsample': lambda module, masks: no_replace(module, masks),
    'LayerNorm': lambda module, masks: replace_layernorm(module, masks),
-    'ConvTranspose2d': lambda module, masks: replace_convtranspose2d(module, masks)
+    'ConvTranspose2d': lambda module, masks: replace_convtranspose2d(module, masks),
+    'Flatten': lambda module, masks: no_replace(module, masks)
 }