Skip to content
Permalink
Browse files

refactor(encoder): no for loop in torch encoder now

  • Loading branch information...
jemmyshin committed Aug 26, 2019
1 parent 37ca2f6 commit 52538276185ffb8c16f10bb70a5f7a1c65aa9076
Showing with 3 additions and 8 deletions.
  1. +3 −8 gnes/encoder/image/torchvision.py
@@ -109,14 +109,9 @@ def _encode(_, img: List['np.ndarray']):
if self._use_cuda:
img_tensor = img_tensor.cuda()

result_npy = []
for t in img_tensor:
t = torch.unsqueeze(t, 0)
encodes = self._model(t)
encodes = torch.squeeze(encodes, 0)
result_npy.append(encodes.data.cpu().numpy())

output = np.array(result_npy, dtype=np.float32)
encodes = self._model(img_tensor)

output = np.array(encodes.data.cpu().numpy(), dtype=np.float32)
return output

output = _encode(self, img)

0 comments on commit 5253827

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