Merge pull request #2215 from beyarkay:patch-1

PiperOrigin-RevId: 556803453

Author: copybara-github

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"
       ]
     },
     {