I think it has something to do with the fact that the grid parameters are created with linspace and meshgrid.
I tried loading in all of the parameters except the grids, but it just gave me an error that the grids were missing.
I tried something along the lines of this solution : How to load part of pre trained model? but again it gave me the same error as the first, that the grids were the same dimensions but it couldn’t copy them.