Update model.py

model.py

@@ -88,9 +88,9 @@ class Generator(nn.Module):
         up6 = self.up6(torch.cat([up5, d3], 1))
         up7 = self.up7(torch.cat([up6, d2], 1))
         return self.final_up(torch.cat([up7, d1], 1))
-def load_model(name):
+# def load_model(name):
     
-    return G.to(device)
+#     return G.to(device)
 
 def gen_model():
     transform = transforms.Compose([
@@ -100,7 +100,6 @@ def gen_model():
     device="cpu"
     # state = torch.load('gen.pth.tar', map_location='cpu')
     # state = state['state_dict']
-    gen=Generator()
+    G=Generator()
     G.load_state_dict(torch.load(f"G101.pth", map_location='cpu'))
-    G.eval()
-    return gen,transform
+    return G,transform