diff --git a/python/caffe/net_spec.py b/python/caffe/net_spec.py index 77a0e0070ae..93fc01927db 100644 --- a/python/caffe/net_spec.py +++ b/python/caffe/net_spec.py @@ -56,8 +56,14 @@ def to_proto(*tops): def assign_proto(proto, name, val): """Assign a Python object to a protobuf message, based on the Python type (in recursive fashion). Lists become repeated fields/messages, dicts - become messages, and other types are assigned directly.""" - + become messages, and other types are assigned directly. For convenience, + repeated fields whose values are not lists are converted to single-element + lists; e.g., `my_repeated_int_field=3` is converted to + `my_repeated_int_field=[3]`.""" + + is_repeated_field = hasattr(getattr(proto, name), 'extend') + if is_repeated_field and not isinstance(val, list): + val = [val] if isinstance(val, list): if isinstance(val[0], dict): for item in val: diff --git a/python/caffe/test/test_net_spec.py b/python/caffe/test/test_net_spec.py index b4595e6531a..fee3c0aaebe 100644 --- a/python/caffe/test/test_net_spec.py +++ b/python/caffe/test/test_net_spec.py @@ -43,8 +43,7 @@ def anon_lenet(batch_size): def silent_net(): n = caffe.NetSpec() - n.data, n.data2 = L.DummyData(shape=[dict(dim=[3]), dict(dim=[4, 2])], - ntop=2) + n.data, n.data2 = L.DummyData(shape=dict(dim=3), ntop=2) n.silence_data = L.Silence(n.data, ntop=0) n.silence_data2 = L.Silence(n.data2, ntop=0) return n.to_proto()