+ encoder = nn.Sequential(
+ Normalizer(mu, std),
+ nn.Conv2d(3, dim_hidden, kernel_size=5, stride=1, padding=2),
+ nn.ReLU(),
+ nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
+ nn.ReLU(),
+ nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
+ nn.ReLU(),
+ nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
+ nn.ReLU(),
+ nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
+ nn.ReLU(),
+ nn.Conv2d(
+ dim_hidden,
+ nb_bits_per_block,
+ kernel_size=block_size,
+ stride=block_size,
+ padding=0,
+ ),
+ SignSTE(),
+ )