Write PyTorch models to cloud locations (requires tensorflow)

io.py

'''
Utilities to save/load torch models.

Same as built-in torch.save/torch.load, but support cloud URLs
'''
import tensorflow.gfile as gio
import io
import torch


def save(object, fname):
    memfile = io.BytesIO()
    torch.save(object, memfile)

    with gio.Open(fname, 'wb') as f:
        f.write(memfile.getvalue())

def load(fname, map_location=None):
    with gio.Open(fname, 'rb') as f:
        memfile = io.BytesIO(f.read())

    return torch.load(memfile, map_location=map_location)