fails in a pytorch script

I am trying to use gradient to run just parts of my code, not the entire script. I have turned the code into a gist which you can find HERE. I want the definition of the model etc to be done on my local machine, then everything after to go on the cloud. I also added a comment with the specific error I get. I initially tried this in a jupyter notebook and it breaks just when you call