Spaces:

timroelofs123
/

face_re-aging_img

Runtime error

timroelofs123 commited on Feb 6, 2024

Commit

4a7c70f

1 Parent(s): a33fac3

fix cuda error

Files changed (1) hide show

app.py CHANGED Viewed

@@ -12,7 +12,7 @@ from scripts.test_functions import process_image
 model_path = "hf/best_unet_model.pth"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 unet_model = UNet().to(device)
-unet_model.load_state_dict(torch.load(model_path))
 unet_model.eval()
 def block(image, source_age, target_age):

 model_path = "hf/best_unet_model.pth"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 unet_model = UNet().to(device)
+unet_model.load_state_dict(torch.load(model_path, map_location=device))
 unet_model.eval()
 def block(image, source_age, target_age):