Unable to load GPU trained model on CPU

[nihit97natu]

I trained this model on google colab on a GPU. I am trying to load the model on my local machine(CPU) using load_model() but unable to load it. I am unable to to use pytorch.save() function too to save the model. It gives some error unable to pickle threads.
How can I solve this?

[metalrt]

Hi @nihit97natu,

I have the same problem. Can you find any solution for this? Thanks

[LiudmylaN]

class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else:
            return super().find_class(module, name)

with open(pathname + '/model.pkl', 'rb') as fp:
        model = CPU_Unpickler(fp).load()