Skip to content
Permalink
Browse files

refactor(onnx): move batch_size to class attribute

  • Loading branch information...
hanxiao committed Aug 2, 2019
1 parent a0fdad3 commit c1e406ae40a53de086873b84669ce20dee0db076
Showing with 2 additions and 3 deletions.
  1. +2 −3 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]

0 comments on commit c1e406a

Please sign in to comment.
You can’t perform that action at this time.