@Sai_Gayatri_Prasad that paper just looks at the trace of FIM, that is much easier to compute. For empirical FIM, it’s the average gradient norm squared. You can compute a batch of gradient norms using the snippet below, and then average them.
A, model = create_toy_model()
activations = {}

def save_activations(layer, a, _):
    activations[layer] = a

with autograd_lib.module_hook(save_activations):
    Y = model(A.t())
    loss = torch.sum(Y * Y) / 2

norms = {}

def compute_norms(layer, _, b):
    a = activations[layer]
    norms[layer] = (a * a).sum(dim=1) * (b * b).sum(dim=1)

with autograd_lib.module_hook(compute_norms):
    loss.backward()

For MNIST MLP, the average gradient norm goes down with time, you won’t see the “inverted V” behavior from the paper

Screenshot 2019-10-28 23.30.12.png960×592 103 KB