import torch


def ridge_regularize(network, lam):
    """Apply ridge penalty at linear layer and hidden-hidden weights."""
    return lam * (
        torch.sum(network.linear.weight ** 2) + torch.sum(network.rnn.weight_hh_l0 ** 2)
    )