(self, x=None)
| 34 | self.batch = batch |
| 35 | |
| 36 | def get_batch(self, x=None): |
| 37 | assert self.batch is not None, "Error: need to set a batch first" |
| 38 | |
| 39 | if x is None or isinstance(self.batch, torch.Tensor): |
| 40 | return self.batch |
| 41 | |
| 42 | # batch is a list; select the corresponding element based on x |
| 43 | size = x.shape[2] |
| 44 | for i in range(len(self.batch)): |
| 45 | if self.batch[i].shape[2] == size: |
| 46 | return self.batch[i] |
| 47 | |
| 48 | raise ValueError("Error: no matching batch found") |
| 49 | |
| 50 | def reset(self): |
| 51 | self.batch = None |