Hi,
No the code looks ok.
The one thing is that you can collapse this:
            ground_truth = ground_truth.to(device)
            ground_truth = ground_truth.to(torch.float32)

into
            ground_truth = ground_truth.to(device=device, dtype=torch.float32)