diff --git a/nbs/01_model_wrapper.ipynb b/nbs/01_model_wrapper.ipynb
index f4f8eee..9717b4b 100644
--- a/nbs/01_model_wrapper.ipynb
+++ b/nbs/01_model_wrapper.ipynb
@@ -55,7 +55,7 @@
     "from fastprogress import master_bar, progress_bar\n",
     "import numpy as np\n",
     "\n",
-    "from torch.tensor import Tensor\n",
+    "from torch import Tensor\n",
     "import torch\n",
     "import torch.nn as nn\n",
     "from torch import optim"
diff --git a/pytorch_inferno/model_wrapper.py b/pytorch_inferno/model_wrapper.py
index 6f9dacd..8af33d0 100644
--- a/pytorch_inferno/model_wrapper.py
+++ b/pytorch_inferno/model_wrapper.py
@@ -12,7 +12,7 @@
 from fastprogress import master_bar, progress_bar
 import numpy as np
 
-from torch.tensor import Tensor
+from torch import Tensor
 import torch
 import torch.nn as nn
 from torch import optim