(self, batch, k)
| 130 | return dec, diff |
| 131 | |
| 132 | def get_input(self, batch, k): |
| 133 | x = batch[k] |
| 134 | if len(x.shape) == 3: |
| 135 | x = x[..., None] |
| 136 | x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() |
| 137 | if self.batch_resize_range is not None: |
| 138 | lower_size = self.batch_resize_range[0] |
| 139 | upper_size = self.batch_resize_range[1] |
| 140 | if self.global_step <= 4: |
| 141 | # do the first few batches with max size to avoid later oom |
| 142 | new_resize = upper_size |
| 143 | else: |
| 144 | new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) |
| 145 | if new_resize != x.shape[2]: |
| 146 | x = F.interpolate(x, size=new_resize, mode="bicubic") |
| 147 | x = x.detach() |
| 148 | return x |
| 149 | |
| 150 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 151 | # https://github.com/pytorch/pytorch/issues/37142 |
no outgoing calls
no test coverage detected