Update app.py

app.py

@@ -136,19 +136,21 @@ print(model)
 model.eval()
 model_2 = Generator(3, 32, 3, 4)
 model_2.load_state_dict(torch.load('G_B_HW4_SAVE.pt',map_location=torch.device('cpu')))
+model_2.eval()
 
 totensor = torchvision.transforms.ToTensor()
 normalize_fn = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 topilimage = torchvision.transforms.ToPILImage()
 
-def predict(input_1, input_2):
-    im1 = normalize_fn(totensor(input))
+def predict(input_1):
+    im1 = normalize_fn(totensor(input_1))
     print(im1.shape)
     preds1 = model(im1.unsqueeze(0))/2 + 0.5
     print(preds1.shape)
-    im2 = normalize_fn(totensor(input))
+
+    im2 = normalize_fn(totensor(input_2))
     print(im2.shape)
-    preds2 = model_2(im1.unsqueeze(0))/2 + 0.5
+    preds2 = model(im2.unsqueeze(0))/2 + 0.5
     print(preds2.shape)
     return topilimage(preds1.squeeze(0).detach()), topilimage(preds2.squeeze(0).detach())