(self)
| 77 | .view(batch_size, 1) |
| 78 | |
| 79 | def get_next_batch(self): |
| 80 | for _ in range(num_batches): |
| 81 | inputs = torch.randn(batch_size, 3, image_w, image_h) |
| 82 | labels = torch.zeros(batch_size, num_classes) \ |
| 83 | .scatter_(1, self.one_hot_indices, 1) |
| 84 | yield inputs.cuda(), labels.cuda() |
| 85 | |
| 86 | def train(self): |
| 87 | name = rpc.get_worker_info().name |