Update app.py

gpu not available so shifting to cpu

app.py

@@ -113,8 +113,7 @@ class SimpleUnet(nn.Module):
         return self.output(x)
 
 model = SimpleUnet()
-print("Num params: ", sum(p.numel() for p in model.parameters()))
-model
+
 
 def linear_beta_schedule(timesteps):
     beta_start = 0.0001
@@ -184,7 +183,7 @@ def sample(model, image_size, batch_size=16, channels=3):
 model = SimpleUnet()
 
 st.title("Generatig images using a diffusion model")
-model.load_state_dict(torch.load("new_linear_model_1090.pt"))
+model.load_state_dict(torch.load("new_linear_model_1090.pt"), map_location=torch.device('cpu'))
 
 
 if(st.button("Click to generate image")):