speeding up FIM computation

Is there a way to speed up FIM computation?

Would perhaps not going through the entire loop in here be ok? e.g. for the extreme case just doing 1 loop
```
for k in range(epochs):
logging.info(f"\tepoch {k + 1}/{epochs}")
for i, (data, target) in enumerate(tqdm(data_loader, leave=False, desc="Computing Fisher")):
data = data.to(device)
output = self.model(data, start_from=self.skip_layers)
# The gradients used to compute the FIM needs to be for y sampled from
# the model distribution y ~ p_w(y|x), not for y from the dataset
if self.bernoulli:
target = torch.bernoulli(F.sigmoid(output)).detach()
else:
target = torch.multinomial(F.softmax(output, dim=-1), 1).detach().view(-1)
loss = self.loss_fn(output, target)
self.model.zero_grad()
loss.backward()
for p in self.model.parameters():
if p.grad is not None:
p.grad2_acc += p.grad.data ** 2
p.grad_counter += 1
break # for debugging faster, otherwise FIM is really slow
break # for debugging faster, otherwise FIM is really slow
```

or are there other better ideas?