diff --git a/lightning_mnist_tpu.ipynb b/lightning_mnist_tpu.ipynb
new file mode 100644
index 00000000..ebff44a8
--- /dev/null
+++ b/lightning_mnist_tpu.ipynb
@@ -0,0 +1,2294 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+ "colab": {
+ "name": "lightning_mnist_tpu.ipynb",
+ "provenance": [],
+ "collapsed_sections": [],
+ "toc_visible": true,
+ "machine_shape": "hm",
+ "include_colab_link": true
+ },
+ "kernelspec": {
+ "name": "python3",
+ "display_name": "Python 3"
+ },
+ "accelerator": "TPU",
+ "widgets": {
+ "application/vnd.jupyter.widget-state+json": {
+ "ed7389ca203749b899be3310fd57efe4": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_2bf5a5c66b5249f6ab3a040942854437",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_cbfd97df2e394c728bc1e1cf39905950",
+ "IPY_MODEL_8628e5720530459fa6cade1eaf5cf197"
+ ]
+ }
+ },
+ "2bf5a5c66b5249f6ab3a040942854437": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "cbfd97df2e394c728bc1e1cf39905950": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_d8e74420bd294df79ec577da56f09bed",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "info",
+ "max": 1,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 1,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_2ab32947d75041218bf52d295944dfb4"
+ }
+ },
+ "8628e5720530459fa6cade1eaf5cf197": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_47c0bd66a884470bbe09de8284c6cc33",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 9920512/? [00:20<00:00, 768075.64it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_fa19c88902c74b63a8175e287e09920c"
+ }
+ },
+ "d8e74420bd294df79ec577da56f09bed": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "2ab32947d75041218bf52d295944dfb4": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "47c0bd66a884470bbe09de8284c6cc33": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "fa19c88902c74b63a8175e287e09920c": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "f5199d5fe4e644019cb29fc0f7978d14": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_2ba65c328f6b416bb450e22780c53e88",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_60711c7201e8407b9cfd81991aeeeda9",
+ "IPY_MODEL_1850df41e3d54040a823ba46faa835a2"
+ ]
+ }
+ },
+ "2ba65c328f6b416bb450e22780c53e88": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "60711c7201e8407b9cfd81991aeeeda9": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_bb6de87fa9294114966d0bc91a74e036",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "success",
+ "max": 1,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 1,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_ece14db2a1854f0291de24ec2c7bdb09"
+ }
+ },
+ "1850df41e3d54040a823ba46faa835a2": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_7689c9eb8c9d45b2847dbd49ef5de720",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 32768/? [00:00<00:00, 89480.58it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_1d1e43b707b44c9f85138e3c8c523809"
+ }
+ },
+ "bb6de87fa9294114966d0bc91a74e036": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "ece14db2a1854f0291de24ec2c7bdb09": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "7689c9eb8c9d45b2847dbd49ef5de720": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "1d1e43b707b44c9f85138e3c8c523809": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "22eb164c77b7400ca202cb77cb85ceba": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_7c9aa99fe36549cc82c31e22d8bb1ec2",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_fbe97940d38b42dc8f814c3969ab10c8",
+ "IPY_MODEL_2de00256ff27477c9fddfcf253c683a9"
+ ]
+ }
+ },
+ "7c9aa99fe36549cc82c31e22d8bb1ec2": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "fbe97940d38b42dc8f814c3969ab10c8": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_3acf1aceab3a4d589e474b0a90913aa9",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "info",
+ "max": 1,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 1,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_541dcc99eabb4ffdbaefc8d1bc8342b0"
+ }
+ },
+ "2de00256ff27477c9fddfcf253c683a9": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_d3752aec0a474043b86b813c5eafbeea",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 1654784/? [00:18<00:00, 375108.62it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_8eef50928fb84726a21d957d83b90a23"
+ }
+ },
+ "3acf1aceab3a4d589e474b0a90913aa9": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "541dcc99eabb4ffdbaefc8d1bc8342b0": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "d3752aec0a474043b86b813c5eafbeea": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "8eef50928fb84726a21d957d83b90a23": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "0162692fcbc7483dbd940e5cffdbfbb0": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_e54c097321fe4d3bb87d5d225232fd74",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_6eec17d4ea9344bb8dc210fbc7538520",
+ "IPY_MODEL_64401a57368d49e4bfba2344fa930e74"
+ ]
+ }
+ },
+ "e54c097321fe4d3bb87d5d225232fd74": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "6eec17d4ea9344bb8dc210fbc7538520": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_fc4a4b7da65e4a689aa2dc9dd2bbddbc",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "success",
+ "max": 1,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 1,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_3ee9362914074eab9aa136f4641703da"
+ }
+ },
+ "64401a57368d49e4bfba2344fa930e74": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_f8e4581a26e940659821d42f98201941",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 8192/? [00:00<00:00, 14903.02it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_2d350214cbab4f7fbc05c9aea5f9c758"
+ }
+ },
+ "fc4a4b7da65e4a689aa2dc9dd2bbddbc": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "3ee9362914074eab9aa136f4641703da": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "f8e4581a26e940659821d42f98201941": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "2d350214cbab4f7fbc05c9aea5f9c758": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "1e63fdda086b49269471fe9f01ce622d": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_defe1c61f3b14b4e974833bbad212942",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_de9c35d40ab04dde8e74dcb830b7c3bf",
+ "IPY_MODEL_91f33c1f3521408da127f2bcda1aec93"
+ ]
+ }
+ },
+ "defe1c61f3b14b4e974833bbad212942": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": "row wrap",
+ "width": "100%",
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": "inline-flex",
+ "left": null
+ }
+ },
+ "de9c35d40ab04dde8e74dcb830b7c3bf": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_ff321181bd694fbea9d6199c702d0dde",
+ "_dom_classes": [],
+ "description": "Validation sanity check: ",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "",
+ "max": 5,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 5,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_67079896993c4e5cb890d0e00e5e2612"
+ }
+ },
+ "91f33c1f3521408da127f2bcda1aec93": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_610da6d3a8304896af37943e3e8bf81f",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 50/? [00:00<00:00, 361.74it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_b63f494872dd49ea97a29e39b8ba3caa"
+ }
+ },
+ "ff321181bd694fbea9d6199c702d0dde": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "67079896993c4e5cb890d0e00e5e2612": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": "2",
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "610da6d3a8304896af37943e3e8bf81f": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "b63f494872dd49ea97a29e39b8ba3caa": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "98cbb9e3816542c79018b7e2c49bba95": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_57b4b6fa412948aa9b345f83048ac00b",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_cb7bdea3309d490081d8fc5e3f2e86d8",
+ "IPY_MODEL_e4877a9015ba44aaa16a54f158a7758b"
+ ]
+ }
+ },
+ "57b4b6fa412948aa9b345f83048ac00b": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": "row wrap",
+ "width": "100%",
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": "inline-flex",
+ "left": null
+ }
+ },
+ "cb7bdea3309d490081d8fc5e3f2e86d8": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_879507c9d60049cdaf2792304fb0893b",
+ "_dom_classes": [],
+ "description": "Epoch 2: 11%",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "danger",
+ "max": 3750,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 400,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_2360bebe316647b7ae59844d14c20c76"
+ }
+ },
+ "e4877a9015ba44aaa16a54f158a7758b": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_8822d3bdee6044e495486761c46ec753",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 400/3750 [00:02<00:17, 195.55it/s, loss=0.812, v_num=0]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_e99a890e1a9e4aea9d165354b4d2b5d1"
+ }
+ },
+ "879507c9d60049cdaf2792304fb0893b": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "2360bebe316647b7ae59844d14c20c76": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": "2",
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "8822d3bdee6044e495486761c46ec753": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "e99a890e1a9e4aea9d165354b4d2b5d1": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "ed0af746e825443b9d64e24c0cd8f150": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HBoxModel",
+ "state": {
+ "_view_name": "HBoxView",
+ "_dom_classes": [],
+ "_model_name": "HBoxModel",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "box_style": "",
+ "layout": "IPY_MODEL_d25e946e31d74b42827b260a8ea87b2c",
+ "_model_module": "@jupyter-widgets/controls",
+ "children": [
+ "IPY_MODEL_c426d275e8b341e485192e092ad0eea0",
+ "IPY_MODEL_a4d4ed5a4eb6435fab748eba79a93c41"
+ ]
+ }
+ },
+ "d25e946e31d74b42827b260a8ea87b2c": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": "row wrap",
+ "width": "100%",
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": "inline-flex",
+ "left": null
+ }
+ },
+ "c426d275e8b341e485192e092ad0eea0": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "FloatProgressModel",
+ "state": {
+ "_view_name": "ProgressView",
+ "style": "IPY_MODEL_6b9279c9e0ea4ae2b31b1105e25cb198",
+ "_dom_classes": [],
+ "description": "Validating: ",
+ "_model_name": "FloatProgressModel",
+ "bar_style": "",
+ "max": 1875,
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": 1875,
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "orientation": "horizontal",
+ "min": 0,
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_0e9a119b77e140058e70c20447df0857"
+ }
+ },
+ "a4d4ed5a4eb6435fab748eba79a93c41": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "HTMLModel",
+ "state": {
+ "_view_name": "HTMLView",
+ "style": "IPY_MODEL_6458d90e5219453bbf8cb3089fc17727",
+ "_dom_classes": [],
+ "description": "",
+ "_model_name": "HTMLModel",
+ "placeholder": "",
+ "_view_module": "@jupyter-widgets/controls",
+ "_model_module_version": "1.5.0",
+ "value": " 1900/? [00:08<00:00, 230.48it/s]",
+ "_view_count": null,
+ "_view_module_version": "1.5.0",
+ "description_tooltip": null,
+ "_model_module": "@jupyter-widgets/controls",
+ "layout": "IPY_MODEL_60ecf9d4f2e745aa99971e5a5f426a1d"
+ }
+ },
+ "6b9279c9e0ea4ae2b31b1105e25cb198": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "ProgressStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "ProgressStyleModel",
+ "description_width": "initial",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "bar_color": null,
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "0e9a119b77e140058e70c20447df0857": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": "2",
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ },
+ "6458d90e5219453bbf8cb3089fc17727": {
+ "model_module": "@jupyter-widgets/controls",
+ "model_name": "DescriptionStyleModel",
+ "state": {
+ "_view_name": "StyleView",
+ "_model_name": "DescriptionStyleModel",
+ "description_width": "",
+ "_view_module": "@jupyter-widgets/base",
+ "_model_module_version": "1.5.0",
+ "_view_count": null,
+ "_view_module_version": "1.2.0",
+ "_model_module": "@jupyter-widgets/controls"
+ }
+ },
+ "60ecf9d4f2e745aa99971e5a5f426a1d": {
+ "model_module": "@jupyter-widgets/base",
+ "model_name": "LayoutModel",
+ "state": {
+ "_view_name": "LayoutView",
+ "grid_template_rows": null,
+ "right": null,
+ "justify_content": null,
+ "_view_module": "@jupyter-widgets/base",
+ "overflow": null,
+ "_model_module_version": "1.2.0",
+ "_view_count": null,
+ "flex_flow": null,
+ "width": null,
+ "min_width": null,
+ "border": null,
+ "align_items": null,
+ "bottom": null,
+ "_model_module": "@jupyter-widgets/base",
+ "top": null,
+ "grid_column": null,
+ "overflow_y": null,
+ "overflow_x": null,
+ "grid_auto_flow": null,
+ "grid_area": null,
+ "grid_template_columns": null,
+ "flex": null,
+ "_model_name": "LayoutModel",
+ "justify_items": null,
+ "grid_row": null,
+ "max_height": null,
+ "align_content": null,
+ "visibility": null,
+ "align_self": null,
+ "height": null,
+ "min_height": null,
+ "padding": null,
+ "grid_auto_rows": null,
+ "grid_gap": null,
+ "max_width": null,
+ "order": null,
+ "_view_module_version": "1.2.0",
+ "grid_template_areas": null,
+ "object_position": null,
+ "object_fit": null,
+ "grid_auto_columns": null,
+ "margin": null,
+ "display": null,
+ "left": null
+ }
+ }
+ }
+ }
+ },
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "view-in-github",
+ "colab_type": "text"
+ },
+ "source": [
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "x83-rnVKT8Wo",
+ "colab_type": "text"
+ },
+ "source": [
+ "# PyTorch Lightning TPU/GPU demo\n",
+ "\n",
+ "On Lightning you can train a model using CPUs, TPUs and GPUs without changing ANYTHING about your code.\n",
+ "\n",
+ "Let's walk through an example!"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "kxtLKs2tan12",
+ "colab_type": "text"
+ },
+ "source": [
+ "---\n",
+ "## Install XLA\n",
+ "First, we have to install the XLA library to support TPUs on PyTorch"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "BHBz1_AnamN_",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "import collections\n",
+ "from datetime import datetime, timedelta\n",
+ "import os\n",
+ "import requests\n",
+ "import threading\n",
+ "\n",
+ "_VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')\n",
+ "VERSION = \"torch_xla==nightly\" #@param [\"xrt==1.15.0\", \"torch_xla==nightly\"]\n",
+ "CONFIG = {\n",
+ " 'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'),\n",
+ " 'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(\n",
+ " (datetime.today() - timedelta(1)).strftime('%Y%m%d'))),\n",
+ "}[VERSION]\n",
+ "DIST_BUCKET = 'gs://tpu-pytorch/wheels'\n",
+ "TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
+ "TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
+ "TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
+ "\n",
+ "# Update TPU XRT version\n",
+ "def update_server_xrt():\n",
+ " print('Updating server-side XRT to {} ...'.format(CONFIG.server))\n",
+ " url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format(\n",
+ " TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0],\n",
+ " XRT_VERSION=CONFIG.server,\n",
+ " )\n",
+ " print('Done updating server-side XRT: {}'.format(requests.post(url)))\n",
+ "\n",
+ "update = threading.Thread(target=update_server_xrt)\n",
+ "update.start()\n",
+ "\n",
+ "# Install Colab TPU compat PyTorch/TPU wheels and dependencies\n",
+ "!pip uninstall -y torch torchvision\n",
+ "!gsutil cp \"$DIST_BUCKET/$TORCH_WHEEL\" .\n",
+ "!gsutil cp \"$DIST_BUCKET/$TORCH_XLA_WHEEL\" .\n",
+ "!gsutil cp \"$DIST_BUCKET/$TORCHVISION_WHEEL\" .\n",
+ "!pip install \"$TORCH_WHEEL\"\n",
+ "!pip install \"$TORCH_XLA_WHEEL\"\n",
+ "!pip install \"$TORCHVISION_WHEEL\"\n",
+ "!sudo apt-get install libomp5\n",
+ "update.join()"
+ ],
+ "execution_count": 0,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "iCsre0XmawoR",
+ "colab_type": "text"
+ },
+ "source": [
+ "## Install PyTorch Lightning\n",
+ "Next, we install PyTorch Lightning"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "dEeUzX_5aLrX",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "! pip install git+git://github.com/williamFalcon/pytorch-lightning.git@master --upgrade"
+ ],
+ "execution_count": 0,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "LXGKqIZuTLMs",
+ "colab_type": "text"
+ },
+ "source": [
+ "## Define lightning model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "aDagOcKyZ_qh",
+ "colab_type": "code",
+ "outputId": "90e8dca3-17b7-4fa9-9dde-54f01017443b",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 71
+ }
+ },
+ "source": [
+ "import os\n",
+ "\n",
+ "import torch\n",
+ "from torch.nn import functional as F\n",
+ "from torch.utils.data import DataLoader\n",
+ "from torchvision.datasets import MNIST\n",
+ "from torchvision import transforms\n",
+ "\n",
+ "import pytorch_lightning as pl\n",
+ "\n",
+ "class CoolSystem(pl.LightningModule):\n",
+ "\n",
+ " def __init__(self):\n",
+ " super(CoolSystem, self).__init__()\n",
+ " # not the best model...\n",
+ " self.l1 = torch.nn.Linear(28 * 28, 10)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " return torch.relu(self.l1(x.view(x.size(0), -1)))\n",
+ "\n",
+ " def training_step(self, batch, batch_idx):\n",
+ " # REQUIRED\n",
+ " x, y = batch\n",
+ " y_hat = self.forward(x)\n",
+ " loss = F.cross_entropy(y_hat, y)\n",
+ " tensorboard_logs = {'train_loss': loss}\n",
+ " return {'loss': loss, 'log': tensorboard_logs}\n",
+ "\n",
+ " def validation_step(self, batch, batch_idx):\n",
+ " # OPTIONAL\n",
+ " x, y = batch\n",
+ " y_hat = self.forward(x)\n",
+ " return {'val_loss': F.cross_entropy(y_hat, y)}\n",
+ "\n",
+ " def validation_end(self, outputs):\n",
+ " # OPTIONAL\n",
+ " avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
+ " tensorboard_logs = {'val_loss': avg_loss}\n",
+ " return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}\n",
+ " \n",
+ " def test_step(self, batch, batch_idx):\n",
+ " # OPTIONAL\n",
+ " x, y = batch\n",
+ " y_hat = self.forward(x)\n",
+ " return {'test_loss': F.cross_entropy(y_hat, y)}\n",
+ "\n",
+ " def test_end(self, outputs):\n",
+ " # OPTIONAL\n",
+ " avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()\n",
+ " tensorboard_logs = {'test_loss': avg_loss}\n",
+ " return {'avg_test_loss': avg_loss, 'log': tensorboard_logs}\n",
+ "\n",
+ " def configure_optimizers(self):\n",
+ " # REQUIRED\n",
+ " # can return multiple optimizers and learning_rate schedulers\n",
+ " # (LBFGS it is automatically supported, no need for closure function)\n",
+ " return torch.optim.Adam(self.parameters(), lr=0.0004)\n",
+ "\n",
+ " def prepare_data(self):\n",
+ " MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())\n",
+ " MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())\n",
+ "\n",
+ " def train_dataloader(self):\n",
+ " # REQUIRED\n",
+ " dataset = MNIST(os.getcwd(), train=True, download=False, transform=transforms.ToTensor())\n",
+ " loader = DataLoader(dataset, batch_size=32)\n",
+ " return loader\n",
+ "\n",
+ " def val_dataloader(self):\n",
+ " dataset = MNIST(os.getcwd(), train=True, download=False, transform=transforms.ToTensor())\n",
+ " loader = DataLoader(dataset, batch_size=32)\n",
+ " return loader\n",
+ "\n",
+ " def test_dataloader(self):\n",
+ " dataset = MNIST(os.getcwd(), train=False, download=False, transform=transforms.ToTensor())\n",
+ " loader = DataLoader(dataset, batch_size=32)\n",
+ " return loader"
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/core/decorators.py:13: UserWarning: data_loader decorator deprecated in 0.6.1. Will remove 0.8.0\n",
+ " warnings.warn(w)\n"
+ ],
+ "name": "stderr"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "yyLkgogITNtj",
+ "colab_type": "text"
+ },
+ "source": [
+ "## Run training"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "kr8cql-aaKnC",
+ "colab_type": "code",
+ "outputId": "0000ac2e-fb15-427d-a27d-5641cbefc5ae",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 634,
+ "referenced_widgets": [
+ "ed7389ca203749b899be3310fd57efe4",
+ "2bf5a5c66b5249f6ab3a040942854437",
+ "cbfd97df2e394c728bc1e1cf39905950",
+ "8628e5720530459fa6cade1eaf5cf197",
+ "d8e74420bd294df79ec577da56f09bed",
+ "2ab32947d75041218bf52d295944dfb4",
+ "47c0bd66a884470bbe09de8284c6cc33",
+ "fa19c88902c74b63a8175e287e09920c",
+ "f5199d5fe4e644019cb29fc0f7978d14",
+ "2ba65c328f6b416bb450e22780c53e88",
+ "60711c7201e8407b9cfd81991aeeeda9",
+ "1850df41e3d54040a823ba46faa835a2",
+ "bb6de87fa9294114966d0bc91a74e036",
+ "ece14db2a1854f0291de24ec2c7bdb09",
+ "7689c9eb8c9d45b2847dbd49ef5de720",
+ "1d1e43b707b44c9f85138e3c8c523809",
+ "22eb164c77b7400ca202cb77cb85ceba",
+ "7c9aa99fe36549cc82c31e22d8bb1ec2",
+ "fbe97940d38b42dc8f814c3969ab10c8",
+ "2de00256ff27477c9fddfcf253c683a9",
+ "3acf1aceab3a4d589e474b0a90913aa9",
+ "541dcc99eabb4ffdbaefc8d1bc8342b0",
+ "d3752aec0a474043b86b813c5eafbeea",
+ "8eef50928fb84726a21d957d83b90a23",
+ "0162692fcbc7483dbd940e5cffdbfbb0",
+ "e54c097321fe4d3bb87d5d225232fd74",
+ "6eec17d4ea9344bb8dc210fbc7538520",
+ "64401a57368d49e4bfba2344fa930e74",
+ "fc4a4b7da65e4a689aa2dc9dd2bbddbc",
+ "3ee9362914074eab9aa136f4641703da",
+ "f8e4581a26e940659821d42f98201941",
+ "2d350214cbab4f7fbc05c9aea5f9c758",
+ "1e63fdda086b49269471fe9f01ce622d",
+ "defe1c61f3b14b4e974833bbad212942",
+ "de9c35d40ab04dde8e74dcb830b7c3bf",
+ "91f33c1f3521408da127f2bcda1aec93",
+ "ff321181bd694fbea9d6199c702d0dde",
+ "67079896993c4e5cb890d0e00e5e2612",
+ "610da6d3a8304896af37943e3e8bf81f",
+ "b63f494872dd49ea97a29e39b8ba3caa",
+ "98cbb9e3816542c79018b7e2c49bba95",
+ "57b4b6fa412948aa9b345f83048ac00b",
+ "cb7bdea3309d490081d8fc5e3f2e86d8",
+ "e4877a9015ba44aaa16a54f158a7758b",
+ "879507c9d60049cdaf2792304fb0893b",
+ "2360bebe316647b7ae59844d14c20c76",
+ "8822d3bdee6044e495486761c46ec753",
+ "e99a890e1a9e4aea9d165354b4d2b5d1",
+ "ed0af746e825443b9d64e24c0cd8f150",
+ "d25e946e31d74b42827b260a8ea87b2c",
+ "c426d275e8b341e485192e092ad0eea0",
+ "a4d4ed5a4eb6435fab748eba79a93c41",
+ "6b9279c9e0ea4ae2b31b1105e25cb198",
+ "0e9a119b77e140058e70c20447df0857",
+ "6458d90e5219453bbf8cb3089fc17727",
+ "60ecf9d4f2e745aa99971e5a5f426a1d"
+ ]
+ }
+ },
+ "source": [
+ "from pytorch_lightning import Trainer\n",
+ "\n",
+ "model = CoolSystem()\n",
+ "\n",
+ "# most basic trainer, uses good defaults\n",
+ "trainer = Trainer(num_tpu_cores=8)\n",
+ "trainer.fit(model) "
+ ],
+ "execution_count": 0,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "INFO:root:GPU available: True, used: True\n",
+ "INFO:root:VISIBLE GPUS: 0\n"
+ ],
+ "name": "stderr"
+ },
+ {
+ "output_type": "display_data",
+ "data": {
+ "text/html": [
+ "
\n",
+ "The default version of TensorFlow in Colab will soon switch to TensorFlow 2.x.
\n",
+ "We recommend you upgrade now \n",
+ "or ensure your notebook will continue to use TensorFlow 1.x via the %tensorflow_version 1.x magic:\n",
+ "more info.