def _register_grad_square_hooks(model: _torch.nn.Module):
        """
        Context manager for registering gradient squaring hooks within the context
        and unregistering them on exit.
        """
        hook_handles = []
        for param in model.parameters():
            if param.requires_grad:
                hook_handles.append(param.register_hook(lambda grad: _torch.square(grad)))
        try:
            yield model
        finally:
            for handle in hook_handles:
                handle.remove()

        Context manager for registering gradient squaring hooks within the context
        and unregistering them on exit.
        
coremltools/optimize/torch/palettization/sensitive_k_means.py