change device in measure computation

mu_map/eval/measures.py

@@ -19,9 +19,10 @@ def nmae(prediction: np.array, target: np.array):
     nmae = mae.sum() / (target.max() - target.min())
     return nmae
 
-device = torch.device("cpu")
+device = torch.device("cuda:0")
 model = UNet()
-model.load_state_dict(torch.load("trainings/01_cgan/snapshots/050_generator.pth", map_location=device))
+model = model.to(device)
+model.load_state_dict(torch.load("trainings/03_cgan/snapshots/50_generator.pth", map_location=device))
 model = model.eval()
 
 transform_normalization = SequenceTransform(transforms=[MeanNormTransform(), PadCropTranform(dim=3, size=32)])
@@ -31,12 +32,9 @@ scores_mse = []
 scores_nmae = []
 for i, (recon, mu_map) in enumerate(dataset):
     print(f"{i:02d}/{len(dataset)}", end="\r")
-    mu_map = crop_or_pad(mu_map)
-    recon = crop_or_pad(recon)
     recon = recon.unsqueeze(dim=0).to(device)
-
-    prediction = model(recon).squeeze().numpy()
-    mu_map = mu_map.squeeze().numpy()
+    prediction = model(recon).squeeze().cpu().numpy()
+    mu_map = mu_map.squeeze().cpu().numpy()
 
     scores_nmae.append(nmae(prediction, mu_map))
     scores_mse.append(mse(prediction, mu_map))