As per the discussion here here with @flo, my problem is now solved. It had to do with package versions in my Julia environment. I needed to downgrade Zygote from v0.4.22 to v0.4.20.