(self, batch, k)
| 243 | return dec, diff |
| 244 | |
| 245 | def get_input(self, batch, k): |
| 246 | x = batch[k] |
| 247 | if len(x.shape) == 3: |
| 248 | x = x[..., None] |
| 249 | x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() |
| 250 | if self.batch_resize_range is not None: |
| 251 | lower_size = self.batch_resize_range[0] |
| 252 | upper_size = self.batch_resize_range[1] |
| 253 | if self.global_step <= 4: |
| 254 | # do the first few batches with max size to avoid later oom |
| 255 | new_resize = upper_size |
| 256 | else: |
| 257 | new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) |
| 258 | if new_resize != x.shape[2]: |
| 259 | x = F.interpolate(x, size=new_resize, mode="bicubic") |
| 260 | x = x.detach() |
| 261 | return x |
| 262 | |
| 263 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 264 | # https://github.com/pytorch/pytorch/issues/37142 |
no outgoing calls
no test coverage detected