def convnet_factory(in_channels, out_channels):
net = nets.ConvResidualNet(
in_channels=in_channels,
out_channels=out_channels,
hidden_channels=hidden_channels,
num_blocks=num_res_blocks,
use_batch_norm=resnet_batchnorm,
dropout_probability=dropout_prob,
)
return net
def create_net(in_channels, out_channels):
# return nets.Conv2d(in_channels, out_channels, kernel_size=1)
return nets.ConvResidualNet(in_channels=in_channels,
out_channels=out_channels,
hidden_channels=16)