Make (simple) wandb logging work again

[relativityhd]

Also auto-formatted the both files with ruff, since I forgot to disable ruff for this project...