Hello Aleco,
I entered the above code into my julia (1.4.1) shell with my package status looking like this
[c52e3926] Atom v0.12.11
[fbb218c0] BSON v0.2.6
[3895d2a7] CUDAapi v4.0.0
[35d6a980] ColorSchemes v3.9.0
[3a865a2d] CuArrays v2.2.0
[31c24e10] Distributions v0.23.2
[ced4e74d] DistributionsAD v0.5.2
[587475ba] Flux v0.10.4
[f67ccb44] HDF5 v0.13.2
[6a3955dd] ImageFiltering v0.6.13
[916415d5] Images v0.22.2
[c8e1da08] IterTools v1.3.0
[033835bb] JLD2 v0.1.13
[e5e0dc1b] Juno v0.8.2
[eb30cadb] MLDatasets v0.5.2
[442fdcdd] Measures v0.3.1
[a3a9e032] NIfTI v0.4.1
[d96e819e] Parameters v0.12.1
[91a5bcdd] Plots v1.3.5
[e88e6eb3] Zygote v0.4.20
[9a3f8284] Random
And the algorithm started printing loss values after calling train() and terminated after three epochs without problems. After what you have written here, I suspect that something is wrong with your julia version or the packages. Can you confirm that you have the same package versions as above? Especially Flux and Zygote should be. You can check by typing ] status into the julia console. Also, when you call julia from your command line does julia 1.4.1 load?