Did you achieve it in fast.ai now or did you stay at pure pytorch? Would be very interested in a code example in fastai