(device, type)
| 55 | |
| 56 | |
| 57 | def _test_permute_batch(device, type): |
| 58 | batch_size = 10 |
| 59 | pipe = Pipeline(batch_size, 4, 0) |
| 60 | data = fn.external_source( |
| 61 | source=lambda: gen_data(batch_size, type), device=device, layout="abc" |
| 62 | ) |
| 63 | perm = fn.batch_permutation() |
| 64 | pipe.set_outputs(data, fn.permute_batch(data, indices=perm), perm) |
| 65 | |
| 66 | for i in range(10): |
| 67 | orig, permuted, idxs = pipe.run() |
| 68 | idxs = [int(idxs.at(i)) for i in range(batch_size)] |
| 69 | orig = orig.as_cpu() |
| 70 | ref = [orig.at(idx) for idx in idxs] |
| 71 | check_batch(permuted, ref, len(ref), 0, 0, "abc") |
| 72 | |
| 73 | |
| 74 | def test_permute_batch(): |
nothing calls this directly
no test coverage detected