(self)
| 26 | self.preload() |
| 27 | |
| 28 | def preload(self): |
| 29 | try: |
| 30 | self.next_input, self.next_target, _, _ = next(self.loader) |
| 31 | except StopIteration: |
| 32 | self.next_input = None |
| 33 | self.next_target = None |
| 34 | return |
| 35 | |
| 36 | with torch.cuda.stream(self.stream): |
| 37 | self.input_cuda() |
| 38 | self.next_target = self.next_target.cuda(non_blocking=True) |
| 39 | |
| 40 | def next(self): |
| 41 | torch.cuda.current_stream().wait_stream(self.stream) |