BTW while our 3D+ML team at AMD is using Julia and AMDGPU.jl, we’re not heavy users of Flux and NNlib yet… we write our ML kernels primarily using KernelAbstractions.jl. So while I wish we could address this Flux+AMDGPU limitation, we can’t prioritize it right now. But our team will be hiring 3 more research engineers soon. If anyone who reads this is interested in joining our team and supporting this use case, please message me. 
12 Likes