move chunk to infer model

2026-09-24 23:33:08 +08:00 · 2022-01-12 02:29:46 +00:00
parent c703a5891a
commit f671f133c0
1 changed files with 7 additions and 4 deletions
@@ -94,8 +94,11 @@ class VQATokenPad(object):
                    'input_ids', 'labels', 'token_type_ids', 'bbox',
                    'attention_mask'
            ]:
-                if self.infer_mode and key == 'labels':
-                    continue
-                length = min(len(data[key]), self.max_seq_len)
-                data[key] = np.array(data[key][:length], dtype='int64')
+                if self.infer_mode:
+                    if key != 'labels':
+                        length = min(len(data[key]), self.max_seq_len)
+                        data[key] = data[key][:length]
+                    else:
+                        continue
+                data[key] = np.array(data[key], dtype='int64')
        return data