model = ResidualFlow(
input_size,
n_blocks=list(map(int, args.nblocks.split('-'))),
intermediate_dim=args.idim,
factor_out=args.factor_out,
quadratic=args.quadratic,
init_layer=init_layer,
actnorm=args.actnorm,
fc_actnorm=args.fc_actnorm,
batchnorm=args.batchnorm,
dropout=args.dropout,
fc=args.fc,
coeff=args.coeff,
vnorms=args.vnorms,
n_lipschitz_iters=args.n_lipschitz_iters,
sn_atol=args.sn_tol,
sn_rtol=args.sn_tol,
n_power_series=args.n_power_series,
n_dist=args.n_dist,
n_samples=args.n_samples,
kernels=args.kernels,
activation_fn=args.act,
fc_end=args.fc_end,
fc_idim=args.fc_idim,
n_exact_terms=args.n_exact_terms,
preact=args.preact,
neumann_grad=args.neumann_grad,
grad_in_forward=args.mem_eff,
first_resblock=args.first_resblock,
learn_p=args.learn_p,
block_type=args.block,
)
model = ResidualFlow(
input_size,
n_blocks=list(map(int, args.nblocks.split('-'))),
intermediate_dim=args.idim,
factor_out=args.factor_out,
quadratic=args.quadratic,
init_layer=init_layer,
actnorm=args.actnorm,
fc_actnorm=args.fc_actnorm,
batchnorm=args.batchnorm,
dropout=args.dropout,
fc=args.fc,
coeff=args.coeff,
vnorms=args.vnorms,
n_lipschitz_iters=args.n_lipschitz_iters,
sn_atol=args.sn_tol,
sn_rtol=args.sn_tol,
n_power_series=args.n_power_series,
n_dist=args.n_dist,
n_samples=args.n_samples,
kernels=args.kernels,
activation_fn=args.act,
fc_end=args.fc_end,
fc_idim=args.fc_idim,
n_exact_terms=args.n_exact_terms,
preact=args.preact,
neumann_grad=args.neumann_grad,
grad_in_forward=args.mem_eff,
first_resblock=args.first_resblock,
learn_p=args.learn_p,
classification=args.task in ['classification', 'hybrid'],
classification_hdim=args.cdim,
n_classes=n_classes,
block_type=args.block,
)
model = ResidualFlow(
input_size,
n_blocks=list(map(int, args.nblocks.split('-'))),
intermediate_dim=args.idim,
factor_out=args.factor_out,
quadratic=args.quadratic,
init_layer=init_layer,
actnorm=args.actnorm,
fc_actnorm=args.fc_actnorm,
batchnorm=args.batchnorm,
dropout=args.dropout,
fc=args.fc,
coeff=args.coeff,
vnorms=args.vnorms,
n_lipschitz_iters=args.n_lipschitz_iters,
sn_atol=args.sn_tol,
sn_rtol=args.sn_tol,
n_power_series=args.n_power_series,
n_dist=args.n_dist,
n_samples=args.n_samples,
kernels=args.kernels,
activation_fn=args.act,
fc_end=args.fc_end,
fc_idim=args.fc_idim,
n_exact_terms=args.n_exact_terms,
preact=args.preact,
neumann_grad=args.neumann_grad,
grad_in_forward=args.mem_eff,
first_resblock=args.first_resblock,
learn_p=args.learn_p,
classification=args.task in ['classification', 'hybrid'],
classification_hdim=args.cdim,
n_classes=n_classes,
block_type=args.block,
)