diff --git a/mindspore/nn/cell.py b/mindspore/nn/cell.py index a694489f5a430605ce1e9739091d7b5801a86fa5..4980e90f3fac9f388d13b68632e4690dba83ed46 100755 --- a/mindspore/nn/cell.py +++ b/mindspore/nn/cell.py @@ -56,7 +56,7 @@ class Cell: >>> def construct(self, x): >>> return self.relu(x) """ - def __init__(self, auto_prefix=True): + def __init__(self, auto_prefix=True, flags=None): self._params = OrderedDict() self._cells = OrderedDict() self.training = False @@ -74,6 +74,8 @@ class Cell: if _get_parallel_mode() in ["auto_parallel", "semi_auto_parallel"]: self._get_construct_inputs_number_and_name() self._parallel_inputs_run = None + if flags: + self.add_flags(**flags) @property def create_time(self): @@ -607,6 +609,11 @@ class Cell: cell.add_flags_recursive(**flags) return self + def get_flags(self): + if not hasattr(self, "_mindspore_flags"): + self._mindspore_flags = {} + return self._mindspore_flags + def to_float(self, dst_type): """ Add cast on all inputs of cell and child cells to run with certain float type. diff --git a/mindspore/nn/wrap/cell_wrapper.py b/mindspore/nn/wrap/cell_wrapper.py index 641558921a7cb4415039b12e110336cd63f4df0a..de0007c2ebe16fcb8be74869356493b0b45c07e2 100644 --- a/mindspore/nn/wrap/cell_wrapper.py +++ b/mindspore/nn/wrap/cell_wrapper.py @@ -219,7 +219,7 @@ class DataWrapper(Cell): """ def __init__(self, network, dataset_types, dataset_shapes, queue_name): - super(DataWrapper, self).__init__(auto_prefix=False) + super(DataWrapper, self).__init__(auto_prefix=False, flags=network.get_flags()) self.get_next = P.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name) self.network = network