diff --git a/test/python/add.py b/test/python/add.py
index 00d4d62..2d37174 100644
--- a/test/python/add.py
+++ b/test/python/add.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import AddNet
 
diff --git a/test/python/gat.py b/test/python/gat.py
index 283c36f..376d682 100644
--- a/test/python/gat.py
+++ b/test/python/gat.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.gat import gat_4_64_8_3
 
diff --git a/test/python/lif.py b/test/python/lif.py
index 7dc797a..f7f3f4c 100644
--- a/test/python/lif.py
+++ b/test/python/lif.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.lif import LIFSumOfSq
 
diff --git a/test/python/mm.py b/test/python/mm.py
index 3c51c37..7625077 100644
--- a/test/python/mm.py
+++ b/test/python/mm.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import MMNet
 
diff --git a/test/python/mul.py b/test/python/mul.py
index fd8692f..34b886d 100644
--- a/test/python/mul.py
+++ b/test/python/mul.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import MulNet
 
diff --git a/test/python/norm.py b/test/python/norm.py
index 27a0848..53017d8 100644
--- a/test/python/norm.py
+++ b/test/python/norm.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import Normalization
 
diff --git a/test/python/resnet.py b/test/python/resnet.py
index 7ac317b..2e5ab5f 100644
--- a/test/python/resnet.py
+++ b/test/python/resnet.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.resnet import resnet_20
 
diff --git a/test/python/scale.py b/test/python/scale.py
index f8a4443..4473782 100644
--- a/test/python/scale.py
+++ b/test/python/scale.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import FeatureScale
 
diff --git a/test/python/spmv.py b/test/python/spmv.py
index 4f52ea0..662e11a 100644
--- a/test/python/spmv.py
+++ b/test/python/spmv.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import MVNet
 
diff --git a/test/python/sqsum.py b/test/python/sqsum.py
index 6f3294c..c96287a 100644
--- a/test/python/sqsum.py
+++ b/test/python/sqsum.py
@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+from mpact.mpactbackend import mpact_jit
 
 from mpact.models.kernels import SqSum
 
diff --git a/test/python/train_simple.py b/test/python/train_simple.py
index dc9d0f6..dd4b0f3 100644
--- a/test/python/train_simple.py
+++ b/test/python/train_simple.py
@@ -96,10 +96,9 @@
 
 # Run it with PyTorch.
 # CHECK-LABEL: pytorch
-# CHECK:       Epoch 99
-# CHECK-SAME:  Accuracy = 1.00
+# CHECK:       Epoch 9
 print("pytorch")
-training_loop(net, optimizer, loss_function, train, validation, epochs=100)
+training_loop(net, optimizer, loss_function, train, validation, epochs=10)
 
 # Run it with MPACT.
 # CHECK-LABEL: mpact