def create_net(in_channels, out_channels):
# return nn.Conv2d(in_channels, out_channels, kernel_size=1)
return nn_.ConvResidualNet(
in_channels=in_channels,
out_channels=out_channels,
hidden_channels=16
)
def create_convnet(in_channels, out_channels):
net = nn_.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