(self, batch: Dict)
| 167 | return loss_mean, loss |
| 168 | |
| 169 | def shared_step(self, batch: Dict) -> Any: |
| 170 | x = self.get_input(batch) |
| 171 | if self.lr_scale is not None: |
| 172 | lr_x = F.interpolate(x, scale_factor=1 / self.lr_scale, mode="bilinear", align_corners=False) |
| 173 | lr_x = F.interpolate(lr_x, scale_factor=self.lr_scale, mode="bilinear", align_corners=False) |
| 174 | lr_z = self.encode_first_stage(lr_x, batch) |
| 175 | batch["lr_input"] = lr_z |
| 176 | |
| 177 | x = x.permute(0, 2, 1, 3, 4).contiguous() |
| 178 | |
| 179 | x = self.encode_first_stage(x, batch) |
| 180 | x = x.permute(0, 2, 1, 3, 4).contiguous() |
| 181 | |
| 182 | gc.collect() |
| 183 | torch.cuda.empty_cache() |
| 184 | |
| 185 | loss, loss_dict = self(x, batch) |
| 186 | |
| 187 | return loss, loss_dict |
| 188 | |
| 189 | def get_input(self, batch): |
| 190 | return batch[self.input_key].to(self.dtype) |
no test coverage detected