pytc.vmc.optimizer.apply_gradient_mask¶
- pytc.vmc.optimizer.apply_gradient_mask(grads, mask)[source]¶
Apply gradient mask to gradients to freeze parameters.
- Parameters:
grads – The gradient PyTree
mask – The mask PyTree created by create_gradient_mask
- Returns:
A PyTree with the same structure as grads, where gradients for frozen parameters are set to zero.