TypeError can't pickle after version update

Hello, getting the “TypeError: can’t pickle X objects” error. As has been posted many times before this is in fact because I am running torch.save(net, PATH) and not torch.save(net.state_dict(), PATH). However, my understanding is that state_dict() just to fill in values so you have to initialize a network with the same architecture to load a state_dict() to it. My application incorporates modifying the architecture after initialization so if that is how it works I can not use state_dict for my saving. Everything was working in pytorch 0.2, I have now updated and working with 0.4. I was wondering if there was any advice on how to track down what exactly is the problem with my object that can’t be pickled. or possibly something that changed between versions that might cause this to no longer work.

If it helps my object is class myObject(torch.autograd.Function). It has a forward and backward and also fills in and keeps track of a few values.

Thanks!