Update pages/23_GANs.py

pages/23_GANs.py

@@ -106,11 +106,12 @@ if train_gan:
 
     # Generate and display images
     z = torch.randn(16, latent_dim)
-    generated_imgs = generator(z).view(-1, 1, 28, 28).data
-    grid = np.transpose(np.array([generated_imgs[i].numpy() for i in range(16)]), (1, 2, 0))
+    generated_imgs = generator(z).view(-1, 1, 28, 28).detach().cpu().numpy()
     
-    fig, ax = plt.subplots(figsize=(8, 8))
-    ax.imshow(np.squeeze(grid), cmap="gray")
+    fig, axes = plt.subplots(4, 4, figsize=(8, 8))
+    for img, ax in zip(generated_imgs, axes.flatten()):
+        ax.imshow(img.reshape(28, 28), cmap="gray")
+        ax.axis('off')
     st.pyplot(fig)
 else:
     st.write("Use the slider to select the number of epochs and click the button to start training the GAN")