(self, x)
| 198 | return log |
| 199 | |
| 200 | def to_rgb(self, x): |
| 201 | assert self.image_key == "segmentation" |
| 202 | if not hasattr(self, "colorize"): |
| 203 | self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) |
| 204 | x = F.conv2d(x, weight=self.colorize) |
| 205 | x = 2.*(x-x.min())/(x.max()-x.min()) - 1. |
| 206 | return x |
| 207 | |
| 208 | |
| 209 | class IdentityFirstStage(torch.nn.Module): |
no test coverage detected