I’ve figured it out
In the testing loop I was calling:
writer.add_histogram("Weights/" + name, param.data.cpu().numpy(), global_step)
writer.add_histogram("Gradients/" + name, param.grad.cpu().numpy(), global_step)

This was calling the 'sort' of 'numpy.ndarray that was slowing down everything