tmabraham commited on
Commit
8d4e13c
1 Parent(s): fc8c230

fix forced bos token, also applying BART model to 8 samples now

Browse files
Files changed (1) hide show
  1. demo/demo_notebook.ipynb +156 -68
demo/demo_notebook.ipynb CHANGED
@@ -11,7 +11,7 @@
11
  },
12
  {
13
  "cell_type": "code",
14
- "execution_count": 1,
15
  "metadata": {
16
  "colab": {
17
  "base_uri": "https://localhost:8080/"
@@ -27,7 +27,7 @@
27
  },
28
  {
29
  "cell_type": "code",
30
- "execution_count": 2,
31
  "metadata": {
32
  "id": "M1wVkrpjU6zO"
33
  },
@@ -39,7 +39,7 @@
39
  },
40
  {
41
  "cell_type": "code",
42
- "execution_count": 3,
43
  "metadata": {},
44
  "outputs": [
45
  {
@@ -65,7 +65,7 @@
65
  },
66
  {
67
  "cell_type": "code",
68
- "execution_count": 4,
69
  "metadata": {
70
  "id": "9jQnM6S2vCpn"
71
  },
@@ -80,7 +80,7 @@
80
  },
81
  {
82
  "cell_type": "code",
83
- "execution_count": 5,
84
  "metadata": {
85
  "id": "_eEaJVxAKpV5"
86
  },
@@ -133,7 +133,7 @@
133
  },
134
  {
135
  "cell_type": "code",
136
- "execution_count": 6,
137
  "metadata": {
138
  "scrolled": true
139
  },
@@ -150,10 +150,10 @@
150
  "text/html": [
151
  "\n",
152
  " Tracking run with wandb version 0.10.33<br/>\n",
153
- " Syncing run <strong style=\"color:#cdcd00\">serene-resonance-1</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
154
  " Project page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax</a><br/>\n",
155
- " Run page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax/runs/1cm35ims\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax/runs/1cm35ims</a><br/>\n",
156
- " Run data is saved locally in <code>/home/tmabraham/vqgan-jax/wandb/run-20210715_030616-1cm35ims</code><br/><br/>\n",
157
  " "
158
  ],
159
  "text/plain": [
@@ -167,20 +167,20 @@
167
  "name": "stderr",
168
  "output_type": "stream",
169
  "text": [
170
- "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-1ef8yxby:v1, 1674.97MB. 2 files... Done. 0:0:0\n"
171
  ]
172
  }
173
  ],
174
  "source": [
175
  "import wandb\n",
176
  "run = wandb.init()\n",
177
- "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-1ef8yxby:v1', type='bart_model')\n",
178
  "artifact_dir = artifact.download()"
179
  ]
180
  },
181
  {
182
  "cell_type": "code",
183
- "execution_count": 7,
184
  "metadata": {
185
  "id": "_6-XKK40oEfP",
186
  "scrolled": true
@@ -203,6 +203,15 @@
203
  "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)"
204
  ]
205
  },
 
 
 
 
 
 
 
 
 
206
  {
207
  "cell_type": "code",
208
  "execution_count": 8,
@@ -253,27 +262,57 @@
253
  {
254
  "cell_type": "code",
255
  "execution_count": 10,
 
 
 
 
 
 
 
 
 
256
  "metadata": {
257
  "id": "P32mJJSbrU1F"
258
  },
259
  "outputs": [],
260
  "source": [
261
- "input_ids_test = tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')"
262
  ]
263
  },
264
  {
265
  "cell_type": "code",
266
- "execution_count": 11,
267
  "metadata": {},
268
  "outputs": [
269
  {
270
  "data": {
271
  "text/plain": [
272
- "DeviceArray([[ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
273
- " 2]], dtype=int32)"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
274
  ]
275
  },
276
- "execution_count": 11,
277
  "metadata": {},
278
  "output_type": "execute_result"
279
  }
@@ -284,18 +323,38 @@
284
  },
285
  {
286
  "cell_type": "code",
287
- "execution_count": 12,
288
  "metadata": {
289
  "id": "C7cHbIHruELT"
290
  },
291
  "outputs": [],
292
  "source": [
293
- "greedy_output = model.generate(input_ids_test, max_length=257)"
294
  ]
295
  },
296
  {
297
  "cell_type": "code",
298
- "execution_count": 13,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
299
  "metadata": {
300
  "colab": {
301
  "base_uri": "https://localhost:8080/"
@@ -307,42 +366,16 @@
307
  {
308
  "data": {
309
  "text/plain": [
310
- "DeviceArray([[16384, 16384, 10042, 10042, 10042, 10042, 10042, 10042,\n",
311
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
312
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
313
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
314
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
315
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
316
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
317
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
318
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
319
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
320
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
321
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
322
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
323
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
324
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
325
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
326
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
327
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
328
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
329
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
330
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
331
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
332
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
333
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
334
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
335
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
336
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
337
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
338
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
339
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
340
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
341
- " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
342
- " 10042]], dtype=int32)"
343
  ]
344
  },
345
- "execution_count": 13,
346
  "metadata": {},
347
  "output_type": "execute_result"
348
  }
@@ -351,6 +384,58 @@
351
  "greedy_output[0]"
352
  ]
353
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
354
  {
355
  "cell_type": "markdown",
356
  "metadata": {},
@@ -360,7 +445,7 @@
360
  },
361
  {
362
  "cell_type": "code",
363
- "execution_count": 14,
364
  "metadata": {},
365
  "outputs": [],
366
  "source": [
@@ -378,7 +463,7 @@
378
  },
379
  {
380
  "cell_type": "code",
381
- "execution_count": 15,
382
  "metadata": {},
383
  "outputs": [],
384
  "source": [
@@ -387,7 +472,7 @@
387
  },
388
  {
389
  "cell_type": "code",
390
- "execution_count": 16,
391
  "metadata": {},
392
  "outputs": [],
393
  "source": [
@@ -402,13 +487,14 @@
402
  },
403
  {
404
  "cell_type": "code",
405
- "execution_count": 17,
406
  "metadata": {
407
  "colab": {
408
  "base_uri": "https://localhost:8080/"
409
  },
410
  "id": "Jz032w73nHEf",
411
- "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49"
 
412
  },
413
  "outputs": [
414
  {
@@ -420,47 +506,49 @@
420
  }
421
  ],
422
  "source": [
423
- "model = VQModel.from_pretrained(\"valhalla/vqgan-imagenet-f16-1024\")"
424
  ]
425
  },
426
  {
427
  "cell_type": "code",
428
- "execution_count": 18,
429
  "metadata": {},
430
  "outputs": [],
431
  "source": [
432
  "def get_images(indices, model):\n",
433
  " indices = indices[:, 1:]\n",
434
- " model.decode_code(indices)\n",
435
- " return indices"
 
436
  ]
437
  },
438
  {
439
  "cell_type": "code",
440
- "execution_count": 19,
441
  "metadata": {},
442
  "outputs": [
443
  {
444
  "name": "stdout",
445
  "output_type": "stream",
446
  "text": [
 
447
  "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
448
  ]
449
  },
450
  {
451
  "data": {
452
- "image/png": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAEACAIAAAD9XIvPAAAAF0lEQVR4nGP4//8/EwMDwygexaN45GEA7ucE/J1FRrMAAAAASUVORK5CYII=\n",
453
  "text/plain": [
454
- "<PIL.Image.Image image mode=RGB size=1x256 at 0x7FE6389B6280>"
455
  ]
456
  },
457
- "execution_count": 19,
458
  "metadata": {},
459
  "output_type": "execute_result"
460
  }
461
  ],
462
  "source": [
463
- "custom_to_pil(np.asarray(get_images(greedy_output[0], model)[0]))"
464
  ]
465
  }
466
  ],
 
11
  },
12
  {
13
  "cell_type": "code",
14
+ "execution_count": null,
15
  "metadata": {
16
  "colab": {
17
  "base_uri": "https://localhost:8080/"
 
27
  },
28
  {
29
  "cell_type": "code",
30
+ "execution_count": 1,
31
  "metadata": {
32
  "id": "M1wVkrpjU6zO"
33
  },
 
39
  },
40
  {
41
  "cell_type": "code",
42
+ "execution_count": 2,
43
  "metadata": {},
44
  "outputs": [
45
  {
 
65
  },
66
  {
67
  "cell_type": "code",
68
+ "execution_count": 3,
69
  "metadata": {
70
  "id": "9jQnM6S2vCpn"
71
  },
 
80
  },
81
  {
82
  "cell_type": "code",
83
+ "execution_count": 4,
84
  "metadata": {
85
  "id": "_eEaJVxAKpV5"
86
  },
 
133
  },
134
  {
135
  "cell_type": "code",
136
+ "execution_count": 5,
137
  "metadata": {
138
  "scrolled": true
139
  },
 
150
  "text/html": [
151
  "\n",
152
  " Tracking run with wandb version 0.10.33<br/>\n",
153
+ " Syncing run <strong style=\"color:#cdcd00\">rare-night-7</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
154
  " Project page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax</a><br/>\n",
155
+ " Run page: <a href=\"https://wandb.ai/tmabraham/vqgan-jax/runs/qzxavce8\" target=\"_blank\">https://wandb.ai/tmabraham/vqgan-jax/runs/qzxavce8</a><br/>\n",
156
+ " Run data is saved locally in <code>/home/tmabraham/vqgan-jax/wandb/run-20210715_075019-qzxavce8</code><br/><br/>\n",
157
  " "
158
  ],
159
  "text/plain": [
 
167
  "name": "stderr",
168
  "output_type": "stream",
169
  "text": [
170
+ "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-1ef8yxby:latest, 1674.97MB. 2 files... Done. 0:0:0\n"
171
  ]
172
  }
173
  ],
174
  "source": [
175
  "import wandb\n",
176
  "run = wandb.init()\n",
177
+ "artifact = run.use_artifact('wandb/hf-flax-dalle-mini/model-1ef8yxby:latest', type='bart_model')\n",
178
  "artifact_dir = artifact.download()"
179
  ]
180
  },
181
  {
182
  "cell_type": "code",
183
+ "execution_count": 6,
184
  "metadata": {
185
  "id": "_6-XKK40oEfP",
186
  "scrolled": true
 
203
  "model = CustomFlaxBartForConditionalGeneration.from_pretrained(artifact_dir)"
204
  ]
205
  },
206
+ {
207
+ "cell_type": "code",
208
+ "execution_count": 7,
209
+ "metadata": {},
210
+ "outputs": [],
211
+ "source": [
212
+ "model.config.forced_bos_token_id = None"
213
+ ]
214
+ },
215
  {
216
  "cell_type": "code",
217
  "execution_count": 8,
 
262
  {
263
  "cell_type": "code",
264
  "execution_count": 10,
265
+ "metadata": {},
266
+ "outputs": [],
267
+ "source": [
268
+ "input_text = ['I enjoy walking with my cute dog']*8"
269
+ ]
270
+ },
271
+ {
272
+ "cell_type": "code",
273
+ "execution_count": 11,
274
  "metadata": {
275
  "id": "P32mJJSbrU1F"
276
  },
277
  "outputs": [],
278
  "source": [
279
+ "input_ids_test = tokenizer(input_text, return_tensors='jax')"
280
  ]
281
  },
282
  {
283
  "cell_type": "code",
284
+ "execution_count": 12,
285
  "metadata": {},
286
  "outputs": [
287
  {
288
  "data": {
289
  "text/plain": [
290
+ "{'input_ids': DeviceArray([[ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
291
+ " 2],\n",
292
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
293
+ " 2],\n",
294
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
295
+ " 2],\n",
296
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
297
+ " 2],\n",
298
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
299
+ " 2],\n",
300
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
301
+ " 2],\n",
302
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
303
+ " 2],\n",
304
+ " [ 0, 100, 2254, 3051, 19, 127, 11962, 2335,\n",
305
+ " 2]], dtype=int32), 'attention_mask': DeviceArray([[1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
306
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
307
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
308
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
309
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
310
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
311
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
312
+ " [1, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=int32)}"
313
  ]
314
  },
315
+ "execution_count": 12,
316
  "metadata": {},
317
  "output_type": "execute_result"
318
  }
 
323
  },
324
  {
325
  "cell_type": "code",
326
+ "execution_count": 13,
327
  "metadata": {
328
  "id": "C7cHbIHruELT"
329
  },
330
  "outputs": [],
331
  "source": [
332
+ "greedy_output = model.generate(input_ids_test['input_ids'], max_length=257)"
333
  ]
334
  },
335
  {
336
  "cell_type": "code",
337
+ "execution_count": 14,
338
+ "metadata": {},
339
+ "outputs": [
340
+ {
341
+ "data": {
342
+ "text/plain": [
343
+ "(8, 257)"
344
+ ]
345
+ },
346
+ "execution_count": 14,
347
+ "metadata": {},
348
+ "output_type": "execute_result"
349
+ }
350
+ ],
351
+ "source": [
352
+ "greedy_output[0].shape"
353
+ ]
354
+ },
355
+ {
356
+ "cell_type": "code",
357
+ "execution_count": 15,
358
  "metadata": {
359
  "colab": {
360
  "base_uri": "https://localhost:8080/"
 
366
  {
367
  "data": {
368
  "text/plain": [
369
+ "DeviceArray([[16384, 10042, 10042, ..., 10042, 10042, 9570],\n",
370
+ " [16384, 10042, 10042, ..., 10042, 10042, 9570],\n",
371
+ " [16384, 10042, 10042, ..., 10042, 10042, 9570],\n",
372
+ " ...,\n",
373
+ " [16384, 10042, 10042, ..., 10042, 10042, 9570],\n",
374
+ " [16384, 10042, 10042, ..., 10042, 10042, 9570],\n",
375
+ " [16384, 10042, 10042, ..., 10042, 10042, 9570]], dtype=int32)"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
376
  ]
377
  },
378
+ "execution_count": 15,
379
  "metadata": {},
380
  "output_type": "execute_result"
381
  }
 
384
  "greedy_output[0]"
385
  ]
386
  },
387
+ {
388
+ "cell_type": "code",
389
+ "execution_count": 16,
390
+ "metadata": {},
391
+ "outputs": [
392
+ {
393
+ "data": {
394
+ "text/plain": [
395
+ "DeviceArray([16384, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
396
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
397
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
398
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
399
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
400
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
401
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
402
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
403
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
404
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
405
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
406
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
407
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
408
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
409
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
410
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
411
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
412
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
413
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
414
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
415
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
416
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
417
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
418
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
419
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
420
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
421
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
422
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
423
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
424
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
425
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
426
+ " 10042, 10042, 10042, 10042, 10042, 10042, 10042, 10042,\n",
427
+ " 9570], dtype=int32)"
428
+ ]
429
+ },
430
+ "execution_count": 16,
431
+ "metadata": {},
432
+ "output_type": "execute_result"
433
+ }
434
+ ],
435
+ "source": [
436
+ "greedy_output[0][0]"
437
+ ]
438
+ },
439
  {
440
  "cell_type": "markdown",
441
  "metadata": {},
 
445
  },
446
  {
447
  "cell_type": "code",
448
+ "execution_count": 17,
449
  "metadata": {},
450
  "outputs": [],
451
  "source": [
 
463
  },
464
  {
465
  "cell_type": "code",
466
+ "execution_count": 18,
467
  "metadata": {},
468
  "outputs": [],
469
  "source": [
 
472
  },
473
  {
474
  "cell_type": "code",
475
+ "execution_count": 19,
476
  "metadata": {},
477
  "outputs": [],
478
  "source": [
 
487
  },
488
  {
489
  "cell_type": "code",
490
+ "execution_count": 20,
491
  "metadata": {
492
  "colab": {
493
  "base_uri": "https://localhost:8080/"
494
  },
495
  "id": "Jz032w73nHEf",
496
+ "outputId": "994d8e85-bff7-480b-8b69-f69dedc15c49",
497
+ "scrolled": true
498
  },
499
  "outputs": [
500
  {
 
506
  }
507
  ],
508
  "source": [
509
+ "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
510
  ]
511
  },
512
  {
513
  "cell_type": "code",
514
+ "execution_count": 21,
515
  "metadata": {},
516
  "outputs": [],
517
  "source": [
518
  "def get_images(indices, model):\n",
519
  " indices = indices[:, 1:]\n",
520
+ " print(indices.shape)\n",
521
+ " img = model.decode_code(indices)\n",
522
+ " return img"
523
  ]
524
  },
525
  {
526
  "cell_type": "code",
527
+ "execution_count": 22,
528
  "metadata": {},
529
  "outputs": [
530
  {
531
  "name": "stdout",
532
  "output_type": "stream",
533
  "text": [
534
+ "(1, 256)\n",
535
  "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
536
  ]
537
  },
538
  {
539
  "data": {
540
+ "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAAtSElEQVR4nO1d6Y4jOXKOyFJ3zwzWXgP2An4Dv/8D+QF8ADvrxczsTFcp/CMzmWRcDDJPqfShu0pi8giScZNS4cevf/3tx28/wjcAJKABBoI7ASEMCEgAAAQAAEhwB0CYgfOD9A4BaHqxVHvhk2NkBZpfU1auvs5bRTqn7GcOEkMDwB1gALgDfcD9HT4Qvt6+f337AkBw/wAkGIaJiREABxhZHgGIgAAw52xBMc7l+GL/F9YjwkWyTmL6USpywRhf3wEGQIK3G+AXgBvcfxre7gQEAG9wJ4Bh4nhEoAEGhDsBIgwA91lSJx1Pmvjh/JOJIAhSoFQM6oRJSDmJOu5KoTG+fCQ7fhqwfUiFcmUS1NWwtghKxgN3YTvVo08uq5kqD2Jgmn8O8A7wFf/4+A6j34MfAIjwls2EAIaxyegOZQPT6DLh/DMrL5ZEEM/oeeEFSPoNAQJakbcUqq14w7Q2ACLQO9Ad7l/hdgP4ABgAhpn1px5nR2jxbUpalkelx4+i5oghLMDPBib90gzmkDXjqtjqp1qfUaJ6z6xnx2JKslXrLUhCwfTyhT6WVTBTMqppoMmrJwQaAEfjcAOCAd5mnwcBhjvcZ84GmmLfRL0VBBclQvcjL9bhro9SOaokzgXZb32vy3na6q5V67cS2UQbaeWqS9AaQEa2HzP2G3mYZneGAG7DXAeBAIjgPpoh4n4cCyqIFiEpplRqOHXuwakF6/T57q06VO2B6TUWgMXjCqZ/QbS1HGAW48k+I1Y34nJU1Y0fxIXgbKdPvYwnncqzWz/mdeBGBDCp/OTxJ7ee8faIez6KJe5ob9TW6Ot4fcgr94vsR63ExDWyNUqwGqtQ1ewdZqlnqYW/nWvXgpuSXNoOU6GSCIhgSLVuOCABEgHNrj5NgewyYen9V9mORJ31HOcgomz9OlKLOEo4SBIYTVRvWHUJZLml8GXnV3cNXdCyCSPy2FLhK8sbz9/SlBTKtTzcgAaEpPAJAHE+85p9MlWXBFwvm6DII9TYsRqQ+XY6TrQq5R0sJZtYy9dKs0PMQ7P+jFz/4qyfze0tuUipM5sIRKCx5hgH3wgIgBCG+7LphHP+NBs170oOzdldUtGxYT47ygrZ66rLq/YktYbqfDNxa4PLnR2Komm0oEnz1QhDt3tbiSvKPbUdHIU9FnORPJEUXtN00gUEcAe4IWJ2wwHnJWEuV04xez2KZrGUpS98ij5q5RXLC3ec745ROshYOVZfh+tjmKaGk5qfi9KdmjocP1MSl7EvzYoebvPYlHqcT3lhdoFSJ6pCV9X9RcCVXJ8WvR4cm9w5P6mKrZgkAtkV639+rQQsKQVv0RlE6nooRqTZIEyHAzcWaNBUCWA6GaiOW6r7uQ+cgumoHO/Dl7zXp+B+EPNYOy3VEWHKVd1Ia4Ols5jblKwJiRdOZw2Q4geFDCxvbwCkTgOLFchluLom468GM6pVsqLH9PQJUh07oXllIltFGlurrapuSc5SYeomsxBsFVcJt5belig53H8dc3cyXFOj7ZyME0ExITzFM9xmxI6khV/HSaTEO8M5SFi/7GONm6/Ip6CA24EtQex38UQdLpLGkd6s5Td3TypivtnoHayZzyhOZHU4KbrcUd/VXWQ5lhq5VPyaHfX4x07U8IaABtsCOOHQfpEk16llNkD4cgttaQF9j9R5ai1lN9dqTmUP1MRDa8PIU2UdrPyGZfjUHXLKzYGVKuR8FkVvUH+CAHAjItTC9CxOzyHdvw1NPNepYg6W0jVCrKKCX265NL5uskasrknVHfDXVlWa3e6WOZZv4/xClgKRXbXQSqlDXMJLtY9pUMttyE91xxW8ZX2g075K3worvy1ynk4oLEbJ7HJrcjcgP+WwrKJ0sZrsieqhOSDxU+2feRkJ0qKCqMNqWjSoU5Yro3erjldjoCQDy1uDNrN9/noAuGX0jBc8qYv7odwSx+IdAC9sYIxjuNi8fv5SZbHAxvuP4krWr8eI5D0sd77q/VdHD07ZaYXygW8cSWyNMZ1CF6SjLUnKTSUqN2SWrXW5+wqmYILB4ma1eIf5W1XTegowNlDVpY60lY8kGV2eyXqYwwY1KPLTKt11txYKAYZR4ZcRBvnNwriEDARirA1AhmCoxAQ7zLulcgg5XIQkB8eskj2ybkwisZS0AMQ7ROXVXHOQn8BBW3p88yQerpagLsiNR/HiLLDFlDxNNcbddU2bZIY1XDEmQ8MupVs85ImNGi9NNdNViNy1IlYplUgBMoiVPnZ6wyQu6iwlw50cMHDlrDk+EoiHpVZzdcRIyeaQc9lcxlY7UVVW0VsRABaBjQUeBoxNbtlbfeLOcgSnSkXvy82LmYl5WMNsljRBunhp41r9SDSxSBX+cCtFqxW5e3PAuNV13gpsXq1e3KjWh1n9V8nuSI8ujWUYR6nTsmb2nLpHlIh4zEFEljgumR2dByGXN95wP7RsKPl82aiw9OfjXVFkypK9SI27l0bVedLf1WhUHbhT/IgjRlrZudQyrM/qDh5gnVoijUpMVA2ZMuhTv6WHVi/yUbdesYYg8ymJn+yRWoKNmx4BCUvGnF454lY0qP2oe0La50gIWq6m28xkRVU+72yC6XPqkcs/xBU36k9mjAKQFs5ZJoKSC6IrmoHtnuTTuBDWBukjMNItK7HeqiWtcDSP2rlTWCEmsMIyA5K/UGVjmy0gl7RqbOMILpSX4aS48LbVHYggX7AgZzc5jrV2lhKFcspHhqlBbEtS4S6R8sjS7pY/q8LUaZbcOLBYuTqY1slUV94GTWth9rZSqKXmzN045sU63BoYxNdAah0ZATkajhGbGq735CNxsj8Km521wLKQ9eCznDWoLNdVaocYtALL3+ztbaQpW/Lq/lV5tZvEDUHiRaRya514YROC6x/spBo7xfuMD6qWs1hujQ+ht03KNNfxaquUfB9yAvlLTinN35mVavLKrYYw9ZL/kw0v6I68sBrLbiucZDTxP2pLS7Von4MsksEOmQ8rbRNNaPA32IZ2LDwgK/fCeWBuSaRBs93wLaAiAJkvNP4qwgyDHXlEJW0eCg8xFwkWylD54oUnxWIBwnyd+yDhAexixwIQrwsA9esZBAuvE2Z8nzpnYSaW/0C8wBYN8cJDoc0ClH+qoiIFbp/Lye/g9KRF6FYIpXtm0vqo+t7qzhnyhaeAqmSd2swBqSQqI2clN6jIk/kBsVk80hzyj84sR5Iqf4PB4mqKziHuhcfHxDL5H9ryPaJSL+t5VCepY8UA1ojM58kTSXr+nCyiXFz28OmF/UGzZz9+ef8BHq8bA5h1y3c1R6Xy2WImPIdM+4UHAO6e+yg6H9TSmZTlpwrTaiwP9ZnITOsr3flCLyzHmZeq+VDvqxGdZulnUt6CCipb5ISaf1/serDo2slcbWIIH96a0mQHgrNgDKaqcpO/vBigQ/dbFcuv2i1infP2Sppayv4B8/eK93sJ7CbcD+VEHuw0hQoPorWpXu4kcvK/BNMJQ+6WYgSwbrSetzPyOI6xe0Fa1Ru8BpasoiaxDyMGcv0Z+jZCuNkE2V2g+upYuR1jaRk/obIpU+x/MAiyuyKWDxagSq2i8d4RoHwxk+7RGOUxbALVeLJjAtmCTAoNY1+PvjRSR41lMJ2mh0XAyyCSM9SUbjtQvNgJOXuTLDKGZ9rn4ibNJG81t0xKiiJBcBCta4mZuqKyvHXok9DHPXnigOlrv39ZgXOwmlsLEMOGuIpUtMTBnSNg5BxghBNG+BWaOkT37eoRdg25WfTpEMN8JMshqVbwHwWnajmDzn5Wt3pDF2t3TahbAOmT+P5PB5Lek1a5W6jKnpyRVy6rdcfE0Vgd67SNY5i6UPfUHd15ZJkOOdqOKrzDQBBnN8MCSJ8EaoO1kELFL3voXpCm7KncHgfVqUhlr8a+rFqrXmxbCb93FD9BkB4e2vLW5NxXgdjv7YHQfhVibbU0MGWN/JBUGyHCTFLfWZobS96ITMXZbCYYB8GfXv420LCuzUjxmoyKa7Gv1+oJgDNyNWTzUY3vGsmJ1GSaWO57Jdff59iu2X8SP3NKWm0K6we1QqMuuAOq7oJTcg4sd6NqAfQVXnkwNHtvfXmuuO/nq+FQP/m29znyKk3WKNKagOBEtcOqxUGjptEkf06yGvJqe0ORvXadyxeycg4gQxtZAbXXEWhaJ9JB64qbtFORaCuq5ROXQbmabWRrxVP0Lk1Be2iVS84OSp1fQXBMhWyjm9zktPKI1WewlWqO+IzqB2G+UlFfdyHYQV9uhDfJGL/4Zm2pa6u77qdufea2DkNkWxKFlgp3UlRNWKPdDIp6+siG5uSso2pCVQAwWUC5H9utkQN1ZGgRmKItc/+xEIYoNUytSeKqpLA+1WokeMeyRYyGDgtQJUwV3ZhWb90yc9B2LNSR0V9zDLCappViYmU7ojWdANl6CoYIQuNqpCEwJnJs0MRwudGwvDIqX0SGaEI15NgB+mgtNCjTbbgLVA0JoEWVnoGF9lzrR2jecEaOM2MVqmRYLlC8gvPI2mW13NcdZcVmiXMatGdieixAW69VlVPui793cmEje21xzvKIRBsIi27cAEUgyaXAGsoww/L7g4UqYU3lENr6cK0ZNrXLVoa70ys2C4DcsEaoIRGbCNo8bSkmKDk5b1s0UR3u+DqqNftEQurDiAMG7hJY1XzD0orVWqBNDLSqRRC8ZlJrLUC7mpH7QvaSSnZ1tlLthJdg9i/fB4zta5DKIILapOpzshL2D9wlzvupVlBfgNiYqv8WRMTrhuiC67WqX4uyPPZrCGUWVGcMvm5gHCufyrdmb8xeRLZH3YnEOpLbNnSWLEhHX65CxMpVg1ppka2AJLbxobWpVlpv0FIQ3BoOFTXcljmPrfFL49AJlqYnz+WZzWr9Btuuh7UTVlSav908OeEMuu1SEH+3hluUJVSyQGLIZpRkJlWLykO9kT8o66FOIWUcP5dgXpJ9v9022EMe1vS5Bz05N6lBmGOsg2BGBnlZT3+UKT5QBaCu80WvVXfELCuGrSJiEnkdNuGsI07PAVr8uSFj+gsuKVOF1SB4mVRVLW+BDs9ZhnycZiEWyhAX3KorQA15R+T63vd/TlpbyxMsMFS+wTBDUxAcznDsCyMY5YnRF/c78HWSFQys93/iiG2fbpBu8gv/O5JxUA9PsFQZR0BGvfLhi/V9tDrD+QvNg2d1j2MG0wKs7zg2iW5O63CKKLWEzEa/NP1OqCY33HZ9IPGzE6sEoHEG3Znx1hkizOyeJ+lTXy8x2BbBBOgaFeg+aWMqRsYQpGsjplEPT5w6/YkEJm0yInlJwinoWHNDwDrVPxOXqAVotG4qWKa/emLYbS6mdG9e9MJK+Oa7a6vW8m7j6KiOOsjWmxxbOBQEEDm4tBQA5s9ean5DVFfyyJA2P8Fv2uJOCyBbdlSI1SXhBckK3tPj8wufAZtwf66YdoLds5732+bzABtE4w1wQh9SXr2wP3o9+5X3GqKntA6G1uaS77bWtexAt3bAMEFx/V84EvGVLy8LrUkObbDbw9xXFNbJX1dTB/n06vmil7t/BQTXP/mo61i4OGvu7ipqATZir3g31WqjvicAovlP0r/wcFixa8G7w/rjVNhsAVag59iibEK5rpep/RceCFvs18JRTm9+FiX6iTALjdOIW6rKLUMqa7y4/xFB7HdzaxIcFeoq55x0EtyZP2mRGys/6ZzSKodiwYP3Fx4C8053b2Y9R2J1PTYbagakPv4O1cVnV7LXL1//GRHaVY17eqJf1QU6BkHuzbW8Mr1gZvSFB8GabNBa3X2wAERQ+fRFcvpfMvAsWOPMovu2jmHl+FtA3tpk09AD3bPJfuEEiE1fmwscAIB9h8+emrX7RsmL258bJ5jzOQgW2I3X2OdSWtu+8MQ4TcEpArAbr/m3OK2aVskLz4Qo023OnMpViCvx2kvxfwqsYLm1HKJchdiN6Viw65xoKZ7SlcTyhY2xguXW8oXyvUC7sVqezMm/qyF/bZIQ+ZDYCw8K6t/Y+jVI//Ewfy/Qkc4GO96Sb81mL5foWdErA2uvKacg+Bjd6ng7ITJeFuApEddrggHWXoQfzmYpNr5+ZfrTfNjlc8xSgMRhVAey9g3LOJzkVFjX3ZRD38/EFJ/Xxeu4Jjm3K+4zo6jl93yFu0AehdYV6mfEp7FzOvrmXg9h/csFDX8mdVPUZ0tqNvSZURXztAzPphCk2rYgEiFUlj/SZbh6+PJ0fL/y5hauj/meFZUvybExWI13RmjMp9tqdvH96ebXC+q/CuFlESOdnhsDPJ2Kr+PF9ArWnQQ3hLwSaz8UvwIVUrf53qOHgRoBqwugfoT6M+ORDsJyeCInd/XZNafl3Fv78nKlRvB1Q4C2c4BtyWlB9VT4M8O3Bms/BnU1HDyBnM/OygJVjPvnFgX11iDU7lA9LuTfqWtsn72OrEhe5wanZYHaNu85tjpDxMVjV2CdRXvaI4IaCADGP3uehcP1RUg1LmUBHg4bzoI0jR69J2vUUft8bGiX4Zy3OnJFcoWrEA+KzW9nb3XIlZj+s1kDH6h+sdpLALohP6JT+UxPDRFtjeKF39uznRxrk1lxmfq8u0A6SGOrC4MZgaadsBx6Z+r5XfFqtc8DylW7O/lizcdXwzU/ZnVBkgSUm9vikZPN9GXcOhOouvWPsHIlWih2UgUT3GVVLkpEPw9w2Lo6bHU24kSpH/9v8o66F+CSK7cLdO+uZgE4lA/Fq/g86+oi7rSol7Q6tH4THtHjbyI459bFHmIoftKfD4gh5t51XS9/rO/nVSySO6Zi9c/+rRzlWlitWzGcKiZZIWoBdkV8Dw/fbcn6Ulod7dMnAzmXOyKH2QtJ26OAoOUvXWvL0TZrLM3xgLFPIx+5tM5Yx+4w0yzlaSNP41jxru+WOIwebCgHfVAxyH61oU3LUKk2rvD16Auuas4tXnRyoMzpV70Xv/MmqKOzkuvss4Lxe4HaFwLZq2oPD3AQdhkxkFyrpnfS21RYvbEje9sQjxgKB4FiZYlZj1Ypv4QASKIfYQMljaQ92jA88OEYmQdA0BOPnAQ3ycAlBEDa70tb6wmRUCXnyDyYXsOmcQvT0fNpCH41olaJH2/5i8tyGpcQgEdDh4Sy3NEaVrO2+MJHiNthKwOXlukSAvBom7bq9tVGPTwf+nibq4PqerFhzvpATAF5tEPli0eGGhi09hBp69dZ38Nl0XydrTgH2JycNei4NHN5RDxSv/nKgOFR0EFqEVAFl0k5B7gaLq+IWg8u0TZme8/18mu5Cj3irZwDXEdLPNd2pYsS7O111vsJMK1n7jZH2iSc9dWIE57R50mI78jeqK7rgy78srbdLKT8lcjjcQUeaYHjl7PTKKnvr3lWdUGSdBB/10y5uA0KcOL81au9D6qOxEVDuAy7X4GGbYDKu+bZ5Q3O+gsxjw5V90juPwwPqzQaISzA2pjqilmga4PKTJrD5f5l6by8aQvVm54vPTai+T7FFQXgIn6DAedCqFqNXAbtMOLqJweeAXRIikw5Cb4arm3O1Vs9KXhR706fxaCO4F3RaKz/Q5EduKIAXDgUZtk2LIlVb0E3Zai3hTPi5bh/E0T+0gwT/eGCBzMXJElA3nOGGuFHTuvy67cC7tzqE+cnwRd0uK9GTwbVZc8DYhRRMmRP4RDWfFTuj9C9njeo3InT/1K8giuRJGlRt0B+6oVK7k8vLuzfnYwQb28Yu0yX4U5Rt5EPU10DHcuDtW3aNQB93FsPBKuJ67jzfU4QfGEnZyWkjyQvwx3D/eSScUFY6YQGROSHB8F9I+2Ey2qnETZ5zjmX6h3tRA6KCo+QUJiwDZ2tonMtAbgeppPa2ke21WW3Tnk30cErD4+fFpR+xHChc4BL7dJMDMJimwnK76WpMTJT9pv7Hn3xSY7rXNjeBvOKR2eEp1iA6uWBE6HRpn4jMSu3gNm/bYGae0PKK73CpbSNhEee49XVZrWos4QTBMA5nT8XyWUpKfGu8ZxO8wjpoQmimRBuGIqcAO1SdAQ80USXigHOtQDhAJer85P4SFJlATN2R1F+BbsrUaGKnHcB6LdBT1cIpxMAAEJBSKIUMg+nXMlpajSoN/Ouj0NJXQTg9BU6nQAAiFGxlQx05ChJiCi436t8jUUNYXTQQwTbLlDbfPEifyAj4Urb5a/LEgqnf+FLWjJy7TE7rMJj5fw1qF98q0BUc9a2NiTADTF1+8irtw3YBZ7WawVEWp5h7i1/bXnwZlbHP0/ALA4uc1aAc+q2ekPjbGxAWgcT3+ZFuQT3n03ExCQ5PzViyjMYpjxxoHpJAQSX1nliTHs7pz/zRPDa3A+u+thnPAC42ofiTyQGl98beIWG/Pg3SZ066hAwf4Yw4vtEHYyzEfLjNlSUy0nwFTTE2RZgxDbeYLaeuXOyeJwdHW6CRz/+ddZBfeQzdvUqxDV48lBsxmnqufJsZOTl0MIXyqQwT3qm+lQm+APUGMcBD4dtb9UmAfjsQfCBk7e8/HwLcuHJ+Z41jFON6pCXROgUbKtZXOgk+DLwT1X9CkE0iRvLt67ChbnfSoXtiwvdBj0RpL40K26yQ92ddDTE3F5cVQaSU3cclNug57pBZ+9N2oMr35NpxXRpjx4jERSycltNhDIB+NQBAAAYYdVWPs+ZmEX50Cz7CgRPP0IQiQSOJABXWZnz6HAOp05E9+g833qVDfawnxY2Z3+tvxBzFsTyXIdbVoUKp3zZ4GpsyY+OI4spBrjmIh1MVW24ay6Sj4NpXs+4eDzFl7MAJxKjrb5HzoXDZJTMdKldtnHcoo4LsqRB5XodL4646iLaBgQIqAdViEYN4+hX6XSLGS5ePrEio/acCNp8Yw/dr9bVszh5LDT/UPYx3K8disprvddBJZGiPUJNnhFXXblrDXApo/yqRqvlFuAmmG6DbnC62Akr53h1tC6XiERTarVpvrxyS+MDUkGtVzPMLoLYimn5QZhwhPZbMjaFZZPOkoPguB3koTY7nMpzeVBlA3FphfF0vqHaNld3G56OR2EpD1lYJYtfhWAOE+6bmpWdE9nP9saaEau3DNCYU+6WpPvSuKgewuWAuln05oAK56yoT2MfcNy1I1Ouhlvf41VWvxVidwtg3eG+LBSpjdavL6awAq3Ht4ven/vB1M9uq4pQfmfeAahlLKw6HOdfhrtOZBY2O6OmKZyZ0g8gyC70+3bZSiiJEiLNOUzqHU29mPrZa5E3jEp7Sewn4WQBmLdtmfspzk9ODxgEZHvDP18rspBosb7oaukBi9dEnGvN0GiWvXTXbTNGb709enzqfD0qArDnFcLcMb0Wcl6EUiYlT1AmNt2iy+xAZk8qmftMYjf+2Lszl32uVZ8jPsuH4o3Bj8kCoSy6CDD7OYNkhezFxL2thzVln4iT9x/Klu7KOGS83Tr70zYJUZsvVJCwShC8rUYxwkGT7tONg0gKT8UkCmn6GwKJK1LQGdoII6nXsABY8x434VRLGLKSVeOsI5LyV8G1q3414jYSXgsHaf5XeLqXMghy72d+z/98Rp7OJ9lqb7h52A26UpOP+WKsSIZuuVQ1v2aBeRt0/51LJn4p2fPYYQNkVlYNCpY4fvw9Ju/dKRFt8zVE9YUj8aIK5lVooxQG+zyLrY8cocd0gbadjKE5uCN0PPszZ4b9y4HKu3EKS6I9fxxRQhvOt7pfOcVxhD3SYvtOUmPKOUCVkuPuAok40s+AHwprERx5SAXlNNrcHlzkpwl8hHgP2JityvPCtvAsxYZB25vF+vs306D7J6UW54HmowA24gFnAh1DqOn2+SijmZu7FnnVzqxZVa2hf6qR3pJdfz30OeXSa8H8btDNT1VUzApjWR2ZYDkAQWPkWImDjZhrmngdiQ5arTSGNo50RZJIVEeO7jnxd53rz4PgI7eR2O8247x2XEvSIhTQ0smY+lATgtHeujB1j/ZJZUpIUbnE60ctB6mWxFNDa6jrbLvEABF7sRKSz0751Db6+korR01NjIofDf2TSjYXa1rWDZ376oFQvhlZc0uNEPAH8ZxkHyWpoKf/W2OjVW65lLH83OQwhydXh4yt1ZL8rXhE2bYv88Ppg2CLFFhX2coOF12Uyw9mP2G+IJ2KpcdddlgM2+ErYLlibsUC5T2a+LC2t7UDWi/DbcalaEz0gMA3h8XxKEpkw1Qx14uJ0bG4Eoc+D+EsNtqj9HO57oYK7bzzPDJJItTHVpp3YC2MOsL23LxVkoZ/K8RhHolYTbV8d0j1prqCqhuDCkvxz8ur5h8Bku8+15HMWR0wElBSOltcv7NysmI47/GBCDHRSN5x3w3K0mDZZuImp6FBGiRnWxV8krotVZ5ec1nEe+gapbyHaY2tGN3qKjyuERnxSKB1qeqy0710CSNNrV+P3i/VdpBNKJTb5rB2wB9RFZj0z/fpd0XVz8jXU11Yw46YcmGH8jLaX34irxPEBgIT7KLyofjNofVfbNAeBFhZmsy9jvYDbqvzjP5CQFWPuLrf9K/sbk1Oa1cH2CswysjBLlq/GrF5UhEfI89+7KFE1Qn6utzBwVGKNv4qEuIxhBxYy37qvXXZ86084YaR9/0LMQ5vzY+qrvCWxDAkXb65i7qWSSuDn2NpcHFWFfsA2YauGKJvdn5k50GmQX1+aOMW1fEo0uP8k6wrvzKtQgyVE6Bi55bCvD6ICmVdzPWWlsdMxm2rxF0+9MHIPdU0r4IlZhnozCzu8clmH5XvBRLocYFQvNY4j+UN9oJ6+JADDcHQuoHc3aTJl8OsnzWBoINtF4mlf+N0qnkN1sPpYdEI02e7MRI3/0gKW4y8d9IVxo4I7katGom9Vx05QqWJ7CpOl0PJGrC5VPtkNKuW1aoc6b+pWrWX3ALzAQBgaGS+bQTaSL1tN8C+yPfe8lnVjd8E1JvBYgkIlrWMQ04Zs3LUCmVlj9BItUAPoV6O/l4gg6BCCeUu08MCA9q9T8n1cQZq7K4WyretSTLsyq6xwTqcMe2NpvvzCqd9MRY576yiB0N1/w62do5iUcM0WCGfq6ZGejqhE1ZeYgS7CyRNmNN2LVmP4O08JXJfBQ2OX7k/608qvB5kuNU91gDF/S21oxejPg38BPdhY+1FRFPzsfIgCn269xKG/pOMFxqQJWx3H2iv5J7GglNZ03gpCxTof1NYVKpZtKcGGbzYnZ8J4hj3c+UU1gTQDWi6CrHjqn1KN8vKonZkOa+MjcW410XRn1gCwBPYG26FGpXLs6UXPhs2OqaM9iRjgHP0Dbsr8cJROCYMgK50fh8r0vS/ZWqKBRCD+yd5azO+qZfazbMXtsWuy0ylXW8ipYOwtrA+lzNFAEipuRfUc4eXHXhkJL5vcGlnQannQaouPpq1insGqYZyF6hRw7/Y9QUGzFRyk/9zXBiYxhgcMtv1//oD8BceFGtu141t6jLjs8d4eyI49mIBqpcu7B6ZMVnLvfnR/OuixMNCXrLo62RNywYJ7LgMR6WNy8s3xl5niS9sj7Wpkdl3z21IUyednBI9CNMub6onOBvjZQc+CcqPATRbj/LDxA3pV08AAjLFxnjp6xfWoJLNdJ61uv4Jg/3I6US1d83KWkxX78HSBi/j8EyI7KaTA9UesU8aohruDnPPFREyxibxswGCsxs+iPTi/mdE966yP7PSMEQ0CD7DuVHGfB2WPSU0PYh2Bb95GwZHdLa6dbQJ0hnJAZfZXzgJ5sbGdjxyDsCfD/nRcTs7r2fFly/zeWHnehrUf1nF/COF1p+w974WJTdFtat9VmRcRfSuyEvrPx/EnpostoWa1APdoW9UNFvtotFfB8PPh5pKjXcS6sDK86jnACguJQWDkHyUXT4H9BKDJ4C7ifzhih2vdDW+9Q7CqmPPQpy+A3BHPwWFEL+uDD0o4lnzwO23pv0vToutD8Wr8DKSWFgA+XxLHDfSC5ui5isn1jzsiyQmBM8BQlE5AlChpmXWXjZlj+rzd06FnYH7UD8gfDTgSZ9BpWwxxQ5ybXZM/n3swhQAFD8h47Pe8fMVkK9hJcv5NyYssv2xq7yS7yvz0JxWDp0quu+6+HVYfSkbKxmNqSStt9YRejhE3frxb4zf2EfF7PHQr8ksWrn91Rbq6yo31tci4ilZYztp3eq+qq5aH46M/jd3LLGiLiO5R+lTBK2EO8yM5e8DrPxAVtYcoWBMawednbXY2vsEgtodszLqUyjrpLcRPUN2TRLjqm6A1QMrDIaMwU2kWGV1xVTKDw7D4jP1O4HcAmSF/d3JLmj57ahUq7MGY6duFWNrEiWyk6abfUkjqb6dZA4pFczCov1I1pH6ILh/tkOiVLNKujuxiiOGolHkLNdj2WT5vUCLQSh+x6Bxf3ChNkbATyvgiERkCMzeolaYykmwL5ZrxKQISmFmP9m/R0A3pQQNQl7UKje3cKgWFwjLx1u5gONvx1Uosd6y6WSQVqKqUlXtWe5TK7lSHqo1H4St4+je4o5W01iOKbtldUk7VQ46wxWoilDjP8dNWiUb8m+fSYXqNufDn2PXTgCbup9EdZzzyvaJZ3rluI9neR+sh5t4VDic2ypky6A3toZsoTcQzrzTarVPwPEMMiZy1lxWzvfIXD+tlNSH8Z0qRyxcTsyqOAdhFBmpA35g4sJxVVBUWEtSR50XBFamWAoXIP53k0oL4FEVOQlmAds2SpcNAO7yUFFRNlW73J7OFzLI0IlEeZPPo/aZl2LVAiMAAWKpu0mxPMtAQ15W1tBdvVMUYSDLQbaZfunuzSFVlpWp0hsbzyzfqclhJtGX1+4GAGHLspCBQLR8+mbJT+2tda3wKpxT0ZoW2gtqk7BCwHhYEq8jyTZZJNZbUxRDYmXUnvXy9uiueE9zJyFrMlXmsWvEYix/JbK67YzdAYCAKNO9p/gcKwYl7QUVc+VQ80AypM8fgegtQnJaUed8Qk3C+L2lPtmOya5IzKINknR72vwJI4vchQAAxu254Nqtpsq39E7bK3X/RguASUbzzTheBpgZBoM1SviqQX2K2iMrnvcfpVFSIbP1mD7Aqn3hK9r80LT8UgaqU9OhRgPxcM13llqoWJCyQDOj6qAkAJQRjdkGwDQfGj89LFaIyv3mThscKxJMA9ekX1/eef/UhUMU3WIpeKICzjzNHuU0LpTMC7ooszjlWaeSxrw/zB9QfbkqjCidpDX+T05YGIsTDsWOWM7VAp4FUn1MVLSUeJk1yc1CUmkHRAiSmDUNjeZKMk66PuE+VcPST/ncvOD47IhTkQRGfNdwaxvlZHXRw/VtUgTV+JYJAKr+DJUywLYrezvqOUpik5bZt0TlaBfHSuY8qM/Z4HDbkpfLUcngxU2h9MocibzQJsPTqGRUI9HhLZcSKoP3oCyKJojaSD5GAUsbk03gUQTjWpD8xEqqMcoO3G+ybIvN9JqO/EOLp4ry64DYHG+AOLIsLgI39gLA00mTYYksjawTaIX59pB8aZa80IZ9VHtkNFRLu1DwQGa4Rq8PUznz6u9LwhQI4IbwcYfhAwaaU/sDwFDyOrtJZqt2Fg7l8dUSaJW2uCCuZvWUmT8jZFTZ9LSPtVgwvKGMqPu1Vf9K5ziJQe7O5N57+iZdRIDb+/udvsAdvhO8Idw/AL8CEsAbAMBAcC850te+ij0qy5OvpdBO9iNjtmzDOgVjN2e3G3GXU3naO53J7Mfv20Q6dGAlqvrGSnnPvOd0CYK4eqX08/brrz/Dn/8F4ctXQIJhALjP/94A3mAAuLM/vyGvFutE9c6mlr60CqtBXUOnmzpah/psazhpNe/n2tKUxPZRVE2nuhGMBnNcnKPN4ZeP39/gh6/w5Q2Gb3O9DwACuAMQwDD/JUnXlqUgYhvIvtS3LRESin+BDszx463Y+kga1CZWBfmoe0abgJHBSrwGgYepw5zxSKvAe8iIQBQ0paQwAtw+/vj7dwAEuAMAwABwm19/zMmcAUbDoBJP4oUzvajpkPVyz5QqnlIQ6vqmEtVByxuy8riaV/VWsEN/CMYKG5od1AjLdXDU83LrSYIV5o4YmfnRdADFSnOWvX376U8Ew/f5Xuj7PJsPgA+Ar3PtAYb7IgOjF7Sk+dNvlaKMXhScpAfD/lpureVURdvRNtjKVNVpt7Q8fQdJlSbZZimikucDkfcmVGonkHGz4b5Y3oA1br56OCf6CQBwFonkAt3+9OOPvwD8CvAV4PsU+wIB/AjwDgAA78V3qKeoYhk7FwPLqOfTFRgPidPDp0/ymEgBaK5T9hxOeymKdiNi1HTIBMwfLjdDbq3pTJZVnJvjHQjGIPj3t69/ALwD/DH5/fdvMADAd4AbwDvAlzkgnnX1EkxjSQ8Zmts2+ZT92oPv8+jfrwOi2ueVwy50iMmYsZnyNrMM1Dk8eyC3LJksmjeeebM0WwCYYoCfhm8E8DPQV0AE+A2GfwDcAb4AfAf4NtsB5AOPfbGSumMbRs6XOZjXVPV0I1YlHsY8J+QSg7a+NaUbhN4NU4LVIMmqgJOfw+ompkmaejokHv7nr//1O8Bv8NuvAB8AAO/fAWD+vpQ3Zf64DMTGjUw0itwC5/8kAaiOXuscWxqu9HSvDnU51qyv0XFDN1SrrXH/cgcnq4bpEQAMQAiIMADgKCe3f/3hT78D/AQ/0eTt3Gjm+3s2GIXUZOFLGPX207J8uVjawlpNEk+0+osMVOfvQPpbfiIj69Nf1VSz6E/7OOxUM5K1sZIytk+7ldbTR6+OJppMEpHMC8F9vvA23OF+A7h9G375PwACGgC/A3yZ4+BvAH8AfJ3dIWEYdQrswuDTbsjkErdQqV7cbbKmF+Eep3PUV5DTVUpIfUCaq2UOYpGpYzqM7WU58hTjseY58YEVONdmImSuDk4q/20mi0Y+v3254z8D/AT48zy9rwA/AgDAN4CP5U9pf5RcBOktOivpwWPCNRym9swcOaNzZyIFRWiUq8BM6gKjm81Dg3EKiwvx4SVdBERVIvlbSVLmoys5eI1StY9WnzYPDRZXaNYdxXEYAbwBfgACwA3ePu4A3wD+AvAd4AeAN4APgC8Av8/WAAFGWSAA9mGL7Bpz6JZoDJStisWO2a8ZKBhd1omMvlF58Zi0QiinxxidtPrt2Mvn9H3i0RzZn01j3cg+LLLlCmHprQMsiXUYH82mbAD4IKA3GN6n4124/UY/fQD8Av/7/f3f4AZ/B/oGCPDHz/D1O/z6T/DtBvcv8EaA75NZZEfC91zgAv7FtZCdkpxrr9dD1QYrPbX1deLwIjHfxyCuR9Kfpxh1NM3uECDgHd4+4GMAeiccEG4//+d//+U//v3Pv3zHb3/72+9/fPvhh/sbfvz6Tl//8eXXjx//9NP7HwBfBrrDgISICESjIMAcV2QTCHqrp4P4u/z7kQr9y4wepVqFmpmL5hcHw+JEy6SwkCB5SjGOJpjOU+t/YK5BRNjyLbTOHDfFJSPjpZQAldWn0GXetjsQIiAMBDQQwv2dAD8A8OPtDd7ubx/w/1lcuBBroJa/AAAAAElFTkSuQmCC\n",
541
  "text/plain": [
542
+ "<PIL.Image.Image image mode=RGB size=256x256 at 0x7FA20677A400>"
543
  ]
544
  },
545
+ "execution_count": 22,
546
  "metadata": {},
547
  "output_type": "execute_result"
548
  }
549
  ],
550
  "source": [
551
+ "custom_to_pil(np.asarray(get_images(jnp.expand_dims(greedy_output[0][0],0), model)[0]))"
552
  ]
553
  }
554
  ],