diff --git a/README.md b/README.md index 342d11f4..a710ca29 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ n_in = nvars + naugs # with augmentation n = 1024 # Model -using ContinuousNormalizingFlows, Lux #, CUDA, ComputationalResources +using ContinuousNormalizingFlows, Lux #, CUDA, ComputationalResources, Zygote nn = Chain(Dense(n_in => 3 * n_in, tanh), Dense(3 * n_in => n_in, tanh)) # icnf = construct(RNODE, nn, nvars) # use defaults icnf = construct( @@ -50,6 +50,8 @@ icnf = construct( naugs; # number of augmented dimensions tspan = (0.0f0, 13.0f0), # have bigger time span steer_rate = 1.0f-1, # add random noise to end of the time span + # compute_mode = DIJacVecVectorMode, # process data one by one + # autodiff_backend = AutoZygote(), # use Zygote # resource = CUDALibs(), # process data by GPU ) @@ -60,7 +62,7 @@ r = rand(data_dist, nvars, n) r = convert.(Float32, r) # Fit It -using DataFrames, MLJBase #, ForwardDiff, ADTypes, OptimizationOptimisers +using DataFrames, MLJBase #, Zygote, ADTypes, OptimizationOptimisers df = DataFrame(transpose(r), :auto) # model = ICNFModel(icnf) # use defaults model = ICNFModel( @@ -68,7 +70,7 @@ model = ICNFModel( batch_size = 256, # have bigger batchs # n_epochs = 100, # have less epochs # optimizers = (Adam(),), # use a different optimizer - # adtype = AutoForwardDiff(), # use ForwardDiff + # adtype = AutoZygote(), # use Zygote ) mach = machine(model, df) fit!(mach)