Update README.md

README.md

@@ -38,6 +38,10 @@ checkpoint = torch.load(checkpoint_path)
 model.load_state_dict(checkpoint['model_state_dict'])
 
 print(f"Validation Accuracy: {checkpoint["val_accuracy"]:.2f}")
+
+input_data = torch.randn(1, 1, 28, 28)
+results = model.predict(input_data.to(device))
+print("Results:", results)
 ```
 
 ## Training Details