Update app.py

app.py

@@ -17,8 +17,8 @@ torch.cuda.empty_cache()
 
 ## Load autoencoder
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
+#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
 autoencoderkl = AutoencoderKL(
     spatial_dims=2,
     in_channels=1,
@@ -33,7 +33,7 @@ autoencoderkl = AutoencoderKL(
 root_dir = "models"
 PATH_auto = f'{root_dir}/auto_encoder_model.pt'
 
-autoencoderkl.load_state_dict(torch.load(PATH_auto))
+autoencoderkl.load_state_dict(torch.load(PATH_auto,map_location=device))
 autoencoderkl = autoencoderkl.to(device)
 
 #### Load unet and embedings
@@ -60,8 +60,8 @@ embed = torch.nn.Embedding(num_embeddings=6, embedding_dim=embedding_dimension,
 PATH_unet_condition = f'{root_dir}/unet_latent_space_model_condition.pt'
 PATH_embed_condition = f'{root_dir}/embed_latent_space_model_condition.pt'
 
-unet.load_state_dict(torch.load(PATH_unet_condition))
-embed.load_state_dict(torch.load(PATH_embed_condition))
+unet.load_state_dict(torch.load(PATH_unet_condition,map_location=device))
+embed.load_state_dict(torch.load(PATH_embed_condition,map_location=device))
 
 # unet.load_state_dict(checkpoint['model_state_dict'])
 # embed.load_state_dict(checkpoint['embed_state_dict'])