Merge pull request #2215 from beyarkay:patch-1

PiperOrigin-RevId: 556803453
tensorflow · Aug 14, 2023 · 007d8f4 · 007d8f4
2 parents 51a06aa + f1da823
commit 007d8f4
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/site/en/tutorials/generative/autoencoder.ipynb b/site/en/tutorials/generative/autoencoder.ipynb
@@ -159,27 +159,29 @@
       },
       "outputs": [],
       "source": [
-        "latent_dim = 64 \n",
-        "\n",
         "class Autoencoder(Model):\n",
-        "  def __init__(self, latent_dim):\n",
+        "  def __init__(self, latent_dim, shape):\n",
         "    super(Autoencoder, self).__init__()\n",
-        "    self.latent_dim = latent_dim   \n",
+        "    self.latent_dim = latent_dim\n",
+        "    self.shape = shape\n",
         "    self.encoder = tf.keras.Sequential([\n",
         "      layers.Flatten(),\n",
         "      layers.Dense(latent_dim, activation='relu'),\n",
         "    ])\n",
         "    self.decoder = tf.keras.Sequential([\n",
-        "      layers.Dense(784, activation='sigmoid'),\n",
-        "      layers.Reshape((28, 28))\n",
+        "      layers.Dense(tf.math.reduce_prod(shape), activation='sigmoid'),\n",
+        "      layers.Reshape(shape)\n",
         "    ])\n",
         "\n",
         "  def call(self, x):\n",
         "    encoded = self.encoder(x)\n",
         "    decoded = self.decoder(encoded)\n",
         "    return decoded\n",
-        "  \n",
-        "autoencoder = Autoencoder(latent_dim) "
+        "\n",
+        "\n",
+        "shape = x_test.shape[1:]\n",
+        "latent_dim = 64\n",
+        "autoencoder = Autoencoder(latent_dim, shape)\n"
       ]
     },
     {