I’m thinking about switching from Flux to Lux, since I really need a WeightNorm layer and I am struggling to write one for myself. I am heavily entrenched into using FluxTraining. Does anyone know if Lux works with FluxTraining (or if there’s a WeightNorm implementation for Flux)?