(images, output)
| 117 | |
| 118 | |
| 119 | def gradient_penalty(images, output): |
| 120 | batch_size = images.shape[0] |
| 121 | |
| 122 | gradients = torch_grad( |
| 123 | outputs=output, |
| 124 | inputs=images, |
| 125 | grad_outputs=torch.ones(output.size(), device=images.device), |
| 126 | create_graph=True, |
| 127 | retain_graph=True, |
| 128 | only_inputs=True, |
| 129 | )[0] |
| 130 | |
| 131 | gradients = rearrange(gradients, "b ... -> b (...)") |
| 132 | return ((gradients.norm(2, dim=1) - 1) ** 2).mean() |
| 133 | |
| 134 | |
| 135 | def leaky_relu(p=0.1): |