Note that any layers using random numbers -- notably `Dropout` and `VariationalHiddenDropout` will give un-expected results till https://github.com/EnzymeAD/Reactant.jl/issues/356