diff --git a/Demo-notebooks/Evaluation.ipynb b/Demo-notebooks/Evaluation.ipynb
index f755b8a..0ccfc44 100644
--- a/Demo-notebooks/Evaluation.ipynb
+++ b/Demo-notebooks/Evaluation.ipynb
@@ -33,7 +33,6 @@
    "source": [
     "# Import libraries\n",
     "import sys\n",
-    "import random\n",
     "import numpy as np\n",
     "import torch\n",
     "import torchio as tio\n",
@@ -43,33 +42,21 @@
     "from src.preprocessing import get_transforms, get_datasets, get_dataloaders, add_vessel_contrast\n",
     "# from src.config import config\n",
     "from src.drr import create_drr\n",
-    "from src.vizualization import plot_results, plot_drr_enhancement\n",
     "import src.training\n",
     "\n",
     "import torch\n",
     "import matplotlib.pyplot as plt\n",
-    "from diffdrr.visualization import plot_drr, plot_mask\n",
     "\n",
     "import torch\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from diffdrr.drr import DRR\n",
     "from diffdrr.data import read, transform_hu_to_density\n",
-    "\n",
-    "from diffdrr.pose import convert\n",
-    "from diffdrr.visualization import drr_to_mesh, img_to_mesh\n",
-    "import pyvista\n",
-    "\n",
-    "from tqdm import tqdm\n",
-    "\n",
     "from src.model import TACEnet\n",
-    "from monai.losses import DiceCELoss\n",
     "from timeit import default_timer as timer\n",
     "import numpy as np\n",
     "import torchvision.transforms as tf\n",
-    "from torchvision import transforms\n",
     "\n",
     "# Set the device\n",
     "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
diff --git a/Demo-notebooks/Training.ipynb b/Demo-notebooks/Training.ipynb
index f7c300e..cbcd47f 100644
--- a/Demo-notebooks/Training.ipynb
+++ b/Demo-notebooks/Training.ipynb
@@ -33,8 +33,6 @@
    "source": [
     "# Import libraries\n",
     "import sys\n",
-    "import random\n",
-    "import numpy as np\n",
     "import torch\n",
     "import torchio as tio\n",
     "\n",
@@ -43,31 +41,10 @@
     "from src.preprocessing import get_transforms, get_datasets, get_dataloaders, add_vessel_contrast\n",
     "# from src.config import config\n",
     "from src.drr import create_drr\n",
-    "from src.vizualization import plot_results, plot_drr_enhancement\n",
+    "from src.vizualization import plot_results\n",
     "import src.training\n",
-    "\n",
-    "import torch\n",
-    "import matplotlib.pyplot as plt\n",
-    "from diffdrr.visualization import plot_drr, plot_mask\n",
-    "\n",
-    "import torch\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "from diffdrr.drr import DRR\n",
     "from diffdrr.data import read, transform_hu_to_density\n",
-    "\n",
-    "from diffdrr.pose import convert\n",
-    "from diffdrr.visualization import drr_to_mesh, img_to_mesh\n",
-    "import pyvista\n",
-    "\n",
-    "from tqdm import tqdm\n",
-    "\n",
     "from src.model import TACEnet\n",
-    "from monai.losses import DiceCELoss\n",
-    "\n",
-    "from torchvision import transforms\n",
     "\n",
     "# Set the device\n",
     "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
diff --git a/Demo-notebooks/demo.ipynb b/Demo-notebooks/demo.ipynb
index 6979b20..6e8a9c8 100644
--- a/Demo-notebooks/demo.ipynb
+++ b/Demo-notebooks/demo.ipynb
@@ -32,7 +32,7 @@
     }
    ],
    "source": [
-    "demo.demonstration(rotation = 35, ef = 0.6, deformation=False)"
+    "demo.demonstration(rotation = 35, ef = 0.6, deformation=False, initial_contrast=4000, device=device)"
    ]
   }
  ],