in vae.py [0:0]
def build(self):
H = self.H
self.in_conv = get_3x3(H.image_channels, H.width)
self.widths = get_width_settings(H.width, H.custom_width_str)
enc_blocks = []
blockstr = parse_layer_string(H.enc_blocks)
for res, down_rate in blockstr:
use_3x3 = res > 2 # Don't use 3x3s for 1x1, 2x2 patches
enc_blocks.append(Block(self.widths[res], int(self.widths[res] * H.bottleneck_multiple), self.widths[res], down_rate=down_rate, residual=True, use_3x3=use_3x3))
n_blocks = len(blockstr)
for b in enc_blocks:
b.c4.weight.data *= np.sqrt(1 / n_blocks)
self.enc_blocks = nn.ModuleList(enc_blocks)