diff --git a/batchflow/models/torch/base_mixins.py b/batchflow/models/torch/base_mixins.py
index 98a791da5..54252f0d7 100644
--- a/batchflow/models/torch/base_mixins.py
+++ b/batchflow/models/torch/base_mixins.py
@@ -6,10 +6,7 @@
 import torch
 
 from ...decorators import deprecated
-
-from ...utils_import import try_import
-plot = try_import(module='...plotter', package=__name__, attribute='plot',
-                  help='Try `pip install batchflow[image]`!')
+from ... import plot
 
 # Also imports `tensorboard`, if necessary
 
diff --git a/pyproject.toml b/pyproject.toml
index 73d1a134e..7805ee920 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,6 @@ dependencies = [
     "llvmlite",
     "scipy>=1.9",
     "tqdm>=4.19",
-    "pytest>=8.3.4",
 ]
 
 [project.optional-dependencies]
@@ -56,7 +55,7 @@ nn = [
 
 dev = [
     "pandas>=0.24",
-    "pytest>=7.0",
+    "pytest>=8.3.4",
     "ruff"
 ]