This is an implementation of Elastic Weight Consolidation (EWC)¹ in PyTorch

import torch.nn.functional as F
from ewc import EWC

module = ...
samples = ...

f = lambda m, s: F.log_softmax(m(s))

ewc = EWC(module, samples, f)

# Code that changes module

l = ewc.loss(module)

The method for calculating the jacobian with respect to the models parameters was based on NNGeometry.

1: James Kirkpatrick, et al. "Overcoming catastrophic forgetting in neural networks". CoRR abs/1612.00796. (2016).