How to add @get_local to a model where the weight is from torch.load() ?

Hi Thanks for sharing your interesting work!
I just met a problem when I load a saved model from torch.load(), the architecture is fixed. So I can't add the @get_local before the definition and 'get_local.cache' returns None.
Do you have any ideas?
