Update README.md

README.md

@@ -21,6 +21,7 @@ library_name: diffusers
     ```python
 
 # Predict function to generate images
+
 def load_model(model_path, device):
     # Initialize the same model architecture as during training
     model = ClassConditionedUnet().to(device)
@@ -60,6 +61,7 @@ def predict(model, class_label, noise_scheduler, num_samples=8, device='cuda'):
     
     return generated_images
 
+
 def display_images(images, num_rows=2):
     # Create a grid of images
     grid = torchvision.utils.make_grid(images, nrow=num_rows)
@@ -72,6 +74,7 @@ def display_images(images, num_rows=2):
     plt.show()
 
 # Example of loading a model and generating predictions
+
 model_path = "model_epoch_0.pth"  # Path to your saved model
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = load_model(model_path, device)