From 4c91f4a2cb8413c90c1d3425e1ef647c75813e96 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Thu, 24 Jul 2025 11:06:24 +0000
Subject: [PATCH] Fix plot for model

---
 batchflow/models/torch/base_mixins.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

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