diff --git a/autoemulate/emulators/neural_net_torch.py b/autoemulate/emulators/neural_net_torch.py
index 9a674e72..bcc3948b 100644
--- a/autoemulate/emulators/neural_net_torch.py
+++ b/autoemulate/emulators/neural_net_torch.py
@@ -1,9 +1,9 @@
 # experimental version of a PyTorch neural network emulator wrapped in Skorch
 # to make it compatible with scikit-learn. Works with cross_validate and GridSearchCV,
 # but doesn't pass tests, because we're subclassing
-import random
+
 import warnings
-from typing import List, Tuple
+from typing import List
 
 import numpy as np
 import torch
@@ -13,7 +13,6 @@
 from skopt.space import Integer, Real
 from skorch import NeuralNetRegressor
 from skorch.callbacks import Callback
-from torch import nn
 
 from autoemulate.emulators.neural_networks import get_module
 from autoemulate.utils import set_random_seed
diff --git a/tests/test_emulators.py b/tests/test_emulators.py
index 41b4230f..30ac5b82 100644
--- a/tests/test_emulators.py
+++ b/tests/test_emulators.py
@@ -152,7 +152,7 @@ def test_nn_sk_pred_type(nn_sk_model, simulation_io):
 
 # Test PyTorch Neural Network (skorch)
 def test_nn_torch_initialisation():
-    nn_torch = NeuralNetTorch()
+    nn_torch = NeuralNetTorch(module="mlp")
     assert nn_torch is not None
 
 
@@ -183,7 +183,9 @@ def test_nn_torch_shape_setter():
     X = np.random.rand(100, input_size)
     y = np.random.rand(100, output_size)
     nn_torch_model = NeuralNetTorch(
-        module__input_size=input_size, module__output_size=output_size
+        module="mlp",
+        module__input_size=input_size,
+        module__output_size=output_size,
     )
     nn_torch_model.fit(X, y)
     assert nn_torch_model.module__input_size == input_size