Update unet/DDPM_Unet_sample.py

unet/DDPM_Unet_sample.py

@@ -911,7 +911,7 @@ save_each = 1
 diffusion_model = diffusion_model.to(device)
 
 last_trained_path = '/content/DDPM_ResNet_Unet/unet/model/epoch_30.pth'
-diffusion_model.load_state_dict(torch.load(os.path.join(last_trained_path))['model'])
+diffusion_model.load_state_dict(torch.load(os.path.join(last_trained_path), map_location=device)['model'])
     
 sample_path = '/content/DDPM_ResNet_Unet/unet/sample'