Update README.md

README.md

@@ -124,9 +124,13 @@ decoder = AttentionDecoder(
     decoder_dim=config["decoder_dim"],
     encoder_dim=config["encoder_dim"],
 ).to(device)
-decoder.load_state_dict(torch.load(decoder_path, map_location=device))
+checkpoint = torch.load(decoder_path, map_location=device)
+decoder.load_state_dict(checkpoint["model_state_dict"])
 decoder.eval()
 
+print(f"Loaded checkpoint from epoch {checkpoint['epoch']}, "
+      f"best val loss: {checkpoint['best_val_loss']:.4f}")
+
 # 5. Beam-search caption generation
 def generate_caption(image_tensor, beam_size=3, max_len=50):
     decoder.eval()