Update README.md

README.md

@@ -62,7 +62,6 @@ from torch import nn
 from huggingface_hub import hf_hub_download
 import matplotlib.pyplot as plt
 
-# 1. Define the Generator class
 class CGAN_Generator(nn.Module):
     def __init__(self, noise_dim=100, num_classes=10, img_size=28):
         super().__init__()
@@ -91,22 +90,18 @@ class CGAN_Generator(nn.Module):
         img = self.conv_blocks(out)
         return img
 
-# 2. Load model weights
-noise_dim = 100
-model = CGAN_Generator(noise_dim=noise_dim)
+model = CGAN_Generator(noise_dim=100)
 
 weights_path = hf_hub_download(repo_id="VioletaR/cgan-mnist", filename="mnist_cgan_generator.pth")
 model.load_state_dict(torch.load(weights_path))
 model.eval()
 
-# 3. Generate a digit (e.g., the number '7')
 target_label = torch.tensor([7]) 
 eval_noise = torch.randn(1, noise_dim)
 
 with torch.no_grad():
     generated_img = model(eval_noise, target_label)
 
-# 4. Denormalize and plot
 generated_img = (generated_img + 1) / 2.0
 img_numpy = generated_img.squeeze().cpu().numpy()