Fix in torchvision2albumentations notebook (albumentations-team#208)

* Fix in torchvision2albumentations notebook
linghai06 · Mar 4, 2019 · 776b252 · 776b252
1 parent 1282f9f
commit 776b252
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/notebooks/migrating_from_torchvision_to_albumentations.ipynb b/notebooks/migrating_from_torchvision_to_albumentations.ipynb
@@ -11,7 +11,8 @@
     "import numpy as np\n",
     "from torch.utils.data import Dataset\n",
     "from torchvision import transforms\n",
-    "from albumentations import Compose, RandomCrop, Normalize, HorizontalFlip, Resize"
+    "from albumentations import Compose, RandomCrop, Normalize, HorizontalFlip, Resize\n",
+    "from albumentations.pytorch import ToTensor"
    ]
   },
   {
@@ -112,7 +113,8 @@
     "    Normalize(\n",
     "        mean=[0.485, 0.456, 0.406],\n",
     "        std=[0.229, 0.224, 0.225],\n",
-    "    )\n",
+    "    ),\n",
+    "    ToTensor()\n",
     "])\n",
     "\n",
     "\n",
@@ -139,7 +141,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -212,6 +214,13 @@
     "| [Resize](https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Resize) \t| [Resize](https://albumentations.readthedocs.io/en/latest/api/augmentations.html#albumentations.augmentations.transforms.Resize) \t| ```Resize(256, 256)``` \t|\n",
     "| [Normalize](https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Normalize) \t| [Normalize](https://albumentations.readthedocs.io/en/latest/api/augmentations.html#albumentations.augmentations.transforms.Normalize) \t| ```Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])``` \t|"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -230,7 +239,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.4"
+   "version": "3.6.1"
   }
  },
  "nbformat": 4,