Solved by EchoShao8899 in post #2 
              
              
                Turn out to be a stupid mistake by me. 
embed_pool = torch.load(os.path.join(args.saved_plugin_dir, 'embed_pool.pth'))

should be changed to 
embed_pool = torch.load(os.path.join(args.saved_plugin_dir, 'embed_pool.pth'), map_location=torch.device('cpu'))

torch.load() will automatically map the file…