(self, batch, k)
| 108 | return dec, posterior |
| 109 | |
| 110 | def get_input(self, batch, k): |
| 111 | x = batch[k] |
| 112 | if len(x.shape) == 3: |
| 113 | x = x[..., None] |
| 114 | x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() |
| 115 | return x |
| 116 | |
| 117 | def training_step(self, batch, batch_idx, optimizer_idx): |
| 118 | inputs = self.get_input(batch, self.image_key) |
no outgoing calls
no test coverage detected