diff --git a/examples/PredictionWithAugmentations.ipynb b/examples/PredictionWithAugmentations.ipynb
index 5a2e88c..e5a66f3 100644
--- a/examples/PredictionWithAugmentations.ipynb
+++ b/examples/PredictionWithAugmentations.ipynb
@@ -216,7 +216,7 @@
     "            prediction = prediction.sigmoid()\n",
     "            \n",
     "            prediction_binary = convert_tensor_to_numpy(prediction.swapaxes(-1, 1))\n",
-    "            transformation.restore_fragment(prediction_binary)\n",
+    "            transformation.restore_fragment(prediction_binary.swapaxes(1,2))\n",
     "\n",
     "tta.merge()"
    ]
diff --git a/requirements.txt b/requirements.txt
index cb0f61f..da47b1e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,5 @@
 PyYAML == 5.3.1
 albumentations == 1.3.0
-numpy == 1.19.1
 opencv_python_headless == 4.4.0.44
 py_oneliner == 0.0.1
 scikit_learn == 0.23.2