| 72 | |
| 73 | # BEGIN run_trainer |
| 74 | def get_next_batch(rank): |
| 75 | for _ in range(10): |
| 76 | num_indices = random.randint(20, 50) |
| 77 | indices = torch.LongTensor(num_indices).random_(0, NUM_EMBEDDINGS) |
| 78 | |
| 79 | # Generate offsets. |
| 80 | offsets = [] |
| 81 | start = 0 |
| 82 | batch_size = 0 |
| 83 | while start < num_indices: |
| 84 | offsets.append(start) |
| 85 | start += random.randint(1, 10) |
| 86 | batch_size += 1 |
| 87 | |
| 88 | offsets_tensor = torch.LongTensor(offsets) |
| 89 | target = torch.LongTensor(batch_size).random_(8).cuda(rank) |
| 90 | yield indices, offsets_tensor, target |
| 91 | |
| 92 | # Train for 100 epochs |
| 93 | for epoch in range(100): |