Repository: pytorch
Issue: [caffe2]Reloading model gives segmentation fault
Labels: caffe2

Here is the code

```python
def create_model(hp_space):
    workspace.ResetWorkspace()
    with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
        train_model = ModelHelper(init_params=True)
        if hp_space['model'] == 'DenseNet':
            obj_train = DenseNet((1, 3, 224, 224), model=train_model,
                                 workspace_name='default', is_test=False)
        else:
            print('Please check hyper parameter name and model name')
            sys.exit(0)
        collect_params = []
        for ops in train_model.InitProto().op:
            for blob_name in ops.output:
                collect_params.append(blob_name)
        workspace.FeedBlob('label', np.ones(4).astype(np.int32))
        obj_train.model.net.Softmax(['fc'], 'softmax')
        xent = obj_train.model.LabelCrossEntropy(['softmax', 'label'],
                                                 'xent')
        loss = obj_train.model.AveragedLoss(xent, 'loss')
        brew.accuracy(obj_train.model, ['softmax', 'label'], 'train')
        train_model.AddGradientOperators([loss])
        if hp_space['optimizer'] == 'adam':
            optimizer.build_adam(train_model,
                                 base_learning_rate=hp_space[
                                     'learning_rate'],
                                 policy='step',
                                 stepsize=1,
                                 gamma=0.999)
        workspace.FeedBlob('input', np.ones((1, 3, 224, 224)).
                           astype(np.float32))
        workspace.RunNetOnce(obj_train.model.param_init_net)
        workspace.CreateNet(obj_train.model.net)
```

Saving using Export method from mobile_exporter

```python
from caffe2.python.predictor.mobile_exporter import Export
def save_checkpoint(model, init_net_name, params, path):
    init_net, predict_net = Export(workspace, model.net, params)
    with open(os.path.join(path, init_net_name), 'wb') as f:
        f.write(init_net.SerializeToString())
```

Loading the model back
```python
def load_trained_model(hp_space, path_to_file):
    workspace.ResetWorkspace()
    with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
        train_model = ModelHelper(init_params=False)
        if hp_space['model'] == 'DenseNet':
            obj_train = DenseNet((1, 3, 224, 224),
                                 model=train_model, workspace_name='default',
                                 is_test=False)
        else:
            print('Please check hyper parameter name and model name')
            sys.exit(0)
        init_net_proto = caffe2_pb2.NetDef()
        with open(os.path.join(path_to_file, 'init_net_deploy.pb'), 'rb') as f:
            init_net_proto.ParseFromString(f.read())
        train_model.param_init_net = train_model.param_init_net.\
            AppendNet(core.Net(init_net_proto))
        collect_params = []
        for ops in train_model.InitProto().op:
            for blob_name in ops.output:
                collect_params.append(blob_name)
        train_model.net.Softmax(['fc'], 'softmax')
        workspace.FeedBlob('label', np.ones(4).astype(np.int32))
        xent = train_model.LabelCrossEntropy(['softmax', 'label'], 'xent')
        loss = train_model.AveragedLoss(xent, 'loss')
        brew.accuracy(train_model, ['softmax', 'label'], 'train')
        train_model.AddGradientOperators([loss])
        if hp_space['optimizer'] == 'adam':
            optimizer.build_adam(train_model,
                                 base_learning_rate=hp_space['learning_rate'],
                                 policy='step',
                                 stepsize=1,
                                 gamma=0.999)
        workspace.RunNetOnce(train_model.param_init_net)
        workspace.CreateNet(train_model.net, overwrite=True)
```

After the model is loaded
```python
workspace.RunNet(train_model.net) # gives segmentation fault
```