AttributeError: Can't get attribute '_VariableFunctionClass' on <module 'torch._C' from 'xxx/torch/_C.cpython-37m-x86_64-linux-gnu.so'>'

I trained a model and stored it with pickle. But when I try to load it on another ubuntu, it shows the error.

AttributeError: Can't get attribute '_VariableFunctionClass' on <module 'torch._C' from 'xxx/torch/_C.cpython-37m-x86_64-linux-gnu.so'>' 

I have tried the method below but it does not work:

    class CustomUnpickler(pickle.Unpickler):

    def find_class(self, module, name):
        try:
            return super().find_class(__name__, name)
        except AttributeError:
            return super().find_class(module, name)

pickle_data = CustomUnpickler(open(path, 'rb')).load()

I also tried with the pickle that are trained with cpu, and it does not trigger that error. So I'm wondering if it's the GPU dependencies that cause this problem.

Thanks in advance!