How to read a pickled collection (list or dictionary etc.) of pytorch cuda tensor in cpu

read_cuda_tensor_in_cpu.md

What if you saved some loss values / accuracy values as a list of pytorch tensor in a system with cuda and then trying to plot the losses in a system with no GPU?

With some googling I found that the following code from (pytorch/pytorch#16797 (comment)) works fine! You just need to define the custome unpickler and use it in place of pickle.load!

import io
import torch
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)

if __name__ == "__main__":
  file_path = "/something/something.pkl"
  file_obj = open(file_path, "rb")
  file_obj.seek(0)
  # instead of pickle.load(file_obj) we do the following
  unpickled_stuff = CPU_Unpickler(file_obj).load()