Use PyTorch UAE in imdb example (#705)

Author: Ashley Scillitoe

doc/source/examples/cd_text_imdb.ipynb

@@ -974,15 +974,14 @@
    "source": [
     "from alibi_detect.cd.pytorch import preprocess_drift\n",
     "from alibi_detect.models.pytorch import TransformerEmbedding\n",
+    "from alibi_detect.cd.pytorch import UAE\n",
     "\n",
+    "# Embedding model\n",
     "embedding_pt = TransformerEmbedding(model_name, emb_type, layers)\n",
     "\n",
-    "model = nn.Sequential(\n",
-    "    embedding_pt,\n",
-    "    nn.Linear(768, 256),\n",
-    "    nn.ReLU(),\n",
-    "    nn.Linear(256, enc_dim)\n",
-    ").to(device).eval()\n",
+    "# PyTorch untrained autoencoder\n",
+    "uae = UAE(input_layer=embedding_pt, shape=shape, enc_dim=enc_dim)\n",
+    "model = uae.to(device).eval()\n",
     "\n",
     "# define preprocessing function\n",
     "preprocess_fn = partial(preprocess_drift, model=model, tokenizer=tokenizer, \n",