fix bug: 垂直方向输入序列长度与水平方向不一样

yizt · Aug 9, 2020 · 66eb4a3 · 66eb4a3
1 parent 721700c
commit 66eb4a3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/generator.py b/generator.py
@@ -190,7 +190,10 @@ def __getitem__(self, item):
 
         target = np.zeros(shape=(self.max_len,), dtype=np.long)
         target[:target_len] = indices
-        input_len = self.im_w // 4 - 3
+        if self.direction == 'horizontal':
+            input_len = self.im_w // 4 - 3
+        else:
+            input_len = self.im_w // 16 - 1
         return image, target, input_len, target_len
 
     def __len__(self):