def METHOD_NAME(self):
    for module_name, ts in self._target_spaces.items():
        wrapper = self._module_wrappers[module_name]
        for target_name, _ in ts.items():
            if hasattr(wrapper, f"{target_name}_scale"):
                delattr(wrapper, f"{target_name}_scale")
            # for deepspeed
            try:
                device = next(wrapper.parameters()).device
            except StopIteration:
                try:
                    device = next(wrapper.buffers()).device
                except StopIteration:
                    # NOTE: this will have risk in model parallel
                    device = next(self.bound_model.parameters()).device
            param = torch.nn.Parameter(torch.Tensor([0.01]).to(device))
            wrapper.register_parameter(f"{target_name}_scale", param)