diff options
| author | Kittipat Virochsiri <kittipat@fb.com> | 2018-02-13 14:18:27 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-02-13 14:40:36 -0800 |
| commit | 83c494787dc9f17692b897b1ea16c2ffdc3f980e (patch) | |
| tree | 98d3572be96ec26b352b0cbac0a83b310d5677ec /caffe2 | |
| parent | 6f533fd8b8732bb38483fcfd0f006b8a0f76052a (diff) | |
| download | pytorch-83c494787dc9f17692b897b1ea16c2ffdc3f980e.tar.gz pytorch-83c494787dc9f17692b897b1ea16c2ffdc3f980e.tar.bz2 pytorch-83c494787dc9f17692b897b1ea16c2ffdc3f980e.zip | |
Allow adding to trainer_extra_schema
Summary: Sometimes we need to add some extra schema later
Reviewed By: sunnieshang
Differential Revision: D6951849
fbshipit-source-id: 564eb88f9250eae24869fd10ba3426e00a18af33
Diffstat (limited to 'caffe2')
| -rw-r--r-- | caffe2/python/layer_model_helper.py | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/caffe2/python/layer_model_helper.py b/caffe2/python/layer_model_helper.py index f20ee64b4f..65f3e0abee 100644 --- a/caffe2/python/layer_model_helper.py +++ b/caffe2/python/layer_model_helper.py @@ -391,6 +391,10 @@ class LayerModelHelper(model_helper.ModelHelper): loss_struct = schema.Struct((prefix, loss)) self._loss = self._loss + loss_struct + def add_trainer_extra_schema(self, trainer_extra_schema): + trainer_extra_record = schema.NewRecord(self.net, trainer_extra_schema) + self._trainer_extra_schema += trainer_extra_record + def __getattr__(self, layer): if layer.startswith('__'): raise AttributeError(layer) |