refactor(onnx): move batch_size to class attribute

gnes-ai · Aug 2, 2019 · c1e406a · c1e406a
1 parent a0fdad3
commit c1e406a
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/gnes/encoder/image/onnx.py b/gnes/encoder/image/onnx.py
@@ -23,15 +23,14 @@
 
 
 class BaseONNXImageEncoder(BaseImageEncoder):
+    batch_size = 64
 
     def __init__(self, model_name: str,
                  model_dir: str,
-                 batch_size: int = 64,
                  use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self.batch_size = batch_size
         self.model_dir = model_dir
         self.model_name = model_name
         self._use_cuda = use_cuda
@@ -66,4 +65,4 @@ def encode(self, img: List['np.ndarray'], *args, **kwargs) -> np.ndarray:
         if pad_batch != 0:
             return result_npy[0][0:len(img)]
         else:
-            return result_npy[0]
+            return result_npy[0]