Update app.py

app.py

@@ -76,6 +76,7 @@ NUM_EPOCHS = 0
 LOAD_MODEL = True
 SAVE_MODEL = False
 CHECKPOINT_GEN_A = f"{os.getcwd()}genA.pth.tar"
+CHECKPOINT_GEN_B = f"{os.getcwd()}genB.pth.tar"
 
 transforms = A.Compose(
     [
@@ -99,10 +100,12 @@ def load_checkpoint(checkpoint_file, model, optimizer, lr):
     for param_group in optimizer.param_groups:
         param_group["lr"] = lr
 
+genB = Generator().to(DEVICE)
 genA = Generator().to(DEVICE)
 optim_gen = optim.Adam(list(genB.parameters()) + list(genA.parameters()), lr=LEARNING_RATE, betas=(0.5, 0.999))
 
 load_checkpoint(CHECKPOINT_GEN_A, genA, optim_gen, LEARNING_RATE)
+load_checkpoint(CHECKPOINT_GEN_B, genB, optim_gen, LEARNING_RATE)
 
 def postprocess_and_show(output):
     # Detach from GPU, move to CPU, and remove the batch dimension