Update classification.py

open-mmlab · Sep 4, 2023 · 79451d8 · 79451d8
1 parent 2e00fca
commit 79451d8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmdeploy/codebase/mmpretrain/deploy/classification.py b/mmdeploy/codebase/mmpretrain/deploy/classification.py
@@ -332,7 +332,8 @@ def get_postprocess(self, *args, **kwargs) -> Dict:
             dict: Composed of the postprocess information.
         """
         postprocess = self.model_cfg.model.head
-        if postprocess['type'] in ('EfficientFormerClsHead', 'StackedLinearClsHead'):
+        if postprocess['type'] in ('EfficientFormerClsHead',
+                                   'StackedLinearClsHead'):
             postprocess['type'] = 'LinearClsHead'
 
         if 'topk' not in postprocess: