From fd9792780d0d969ab95463975842bc24dfba816c Mon Sep 17 00:00:00 2001
From: zhang chen <chen.zhang@zilliz.com>
Date: Fri, 17 Dec 2021 17:24:30 +0800
Subject: [PATCH] add return shape introduction

---
 README.md        | 1 +
 pytorch/model.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 4ffc7d1..be4df3b 100644
--- a/README.md
+++ b/README.md
@@ -40,6 +40,7 @@ The Operator returns a tuple `Tuple[('feature_vector', numpy.ndarray)]` containi
 - feature_vector:
   - the embedding of the image
   - data type: `numpy.ndarray`
+  - shape: (dim,)
 
 ## Requirements
 
diff --git a/pytorch/model.py b/pytorch/model.py
index 5a3a050..4a5fbc8 100644
--- a/pytorch/model.py
+++ b/pytorch/model.py
@@ -32,7 +32,7 @@ class Model():
         self._model.eval()
 
     def __call__(self, img_tensor: torch.Tensor):
-        return self._model(img_tensor).detach().numpy()
+        return self._model(img_tensor).flatten().detach().numpy()
 
     def train(self):
         """