Just a thought, but have you considered simply using PyTorch through PyCall.jl or something like ThArrays.jl? This may not be worth it if the majority of your code is reliant on PyTorch, but it's certainly an option that should be taken seriously imo.