(self)
| 109 | |
| 110 | class Encoder: |
| 111 | def __init__(self): |
| 112 | sz = [(128, 128), (128, 256), (256, 512), (512, 512)] |
| 113 | self.conv_in = Conv2d(3,128,3, padding=1) |
| 114 | |
| 115 | arr = [] |
| 116 | for i,s in enumerate(sz): |
| 117 | arr.append({"block": |
| 118 | [ResnetBlock(s[0], s[1]), |
| 119 | ResnetBlock(s[1], s[1])]}) |
| 120 | if i != 3: arr[-1]['downsample'] = {"conv": Conv2d(s[1], s[1], 3, stride=2, padding=(0,1,0,1))} |
| 121 | self.down = arr |
| 122 | |
| 123 | self.mid = Mid(512) |
| 124 | self.norm_out = GroupNorm(32, 512) |
| 125 | self.conv_out = Conv2d(512, 8, 3, padding=1) |
| 126 | |
| 127 | def __call__(self, x): |
| 128 | x = self.conv_in(x) |
nothing calls this directly
no test coverage detected