瀏覽代碼

Support Cython generator type (#2136)

Lin Manhui 1 年之前
父節點
當前提交
de44165d96
共有 1 個文件被更改,包括 4 次插入1 次删除
  1. 4 1
      paddlex/inference/components/base.py

+ 4 - 1
paddlex/inference/components/base.py

@@ -49,7 +49,10 @@ class BaseComponent(ABC):
                         yield item
 
             # output may be a generator, when the apply() uses yield
-            if isinstance(output, GeneratorType):
+            if (
+                isinstance(output, GeneratorType)
+                or output.__class__.__name__ == "generator"
+            ):
                 # if output is a generator, use for-in to get every one batch output data and yield one by one
                 for each_output in output:
                     reassemble_data = self._check_output(each_output, input_)