def update_data_attribute(
        self: TAvalancheDataset, name: str, new_value
    ) -> TAvalancheDataset:
        """
        Return a new dataset with the added or replaced data attribute.

        If a object of type :class:`DataAttribute` is passed, then the data
        attribute is setted as is.

        Otherwise, if a raw value is passed, a new DataAttribute is created.
        If a DataAttribute with the same already exists, the use_in_getitem
        flag is inherited, otherwise it is set to False.

        :param name: The name of the data attribute to add/replace.
        :param new_value: Either a :class:`DataAttribute` or a sequence
            containing as many elements as the datasets.
        :returns: A copy of this dataset with the given data attribute set.
        """
        assert len(new_value) == len(
            self
        ), f"Size mismatch when updating data attribute {name}"

        datacopy = self._shallow_clone_dataset()
        datacopy._data_attributes = copy.copy(datacopy._data_attributes)

        if isinstance(new_value, DataAttribute):
            assert name == new_value.name
            datacopy._data_attributes[name] = new_value
        else:
            use_in_getitem = False
            prev_attr = datacopy._data_attributes.get(name, None)
            if prev_attr is not None:
                use_in_getitem = prev_attr.use_in_getitem

            datacopy._data_attributes[name] = DataAttribute(
                new_value, name=name, use_in_getitem=use_in_getitem
            )

        if not hasattr(datacopy, name):
            # Creates the field if it does not exist
            setattr(datacopy, name, datacopy._data_attributes[name])

        return datacopy

        Return a new dataset with the added or replaced data attribute.

        If a object of type :class:`DataAttribute` is passed, then the data
        attribute is setted as is.

        Otherwise, if a raw value is passed, a new DataAttribute is created.
        If a DataAttribute with the same already exists, the use_in_getitem
        flag is inherited, otherwise it is set to False.

        :param name: The name of the data attribute to add/replace.
        :param new_value: Either a :class:`DataAttribute` or a sequence
            containing as many elements as the datasets.
        :returns: A copy of this dataset with the given data attribute set.
        
avalanche/benchmarks/utils/data.py