| 32 | |
| 33 | |
| 34 | class TMPCollate: |
| 35 | def __call__(self, input_list): |
| 36 | batch_size = len(input_list) |
| 37 | seq_lens = [s.shape[0] for s in input_list] |
| 38 | dim = input_list[0].shape[1] |
| 39 | out = np.zeros((batch_size, max(seq_lens), dim), dtype='int32') # pad with 0, int32 to be free of precision |
| 40 | for i in range(batch_size): |
| 41 | out[i, :seq_lens[i]] = input_list[i] |
| 42 | return np.concatenate(input_list, axis=0) |
| 43 | |
| 44 | |
| 45 | class MyTestCase(unittest.TestCase): |