fix missing model and taskmodule imports from pie_modules in evaluate…

….py and predict.py
ArneBinder · Jun 12, 2024 · 1d3d31a · 1d3d31a
1 parent 76b0454
commit 1d3d31a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/evaluate.py b/src/evaluate.py
@@ -39,6 +39,8 @@
 import pytorch_lightning as pl
 from omegaconf import DictConfig
 from pie_datasets import DatasetDict
+from pie_modules.models import *  # noqa: F403
+from pie_modules.taskmodules import *  # noqa: F403
 from pytorch_ie.core import PyTorchIEModel, TaskModule
 from pytorch_lightning import Trainer
 

diff --git a/src/predict.py b/src/predict.py
@@ -40,6 +40,8 @@
 import pytorch_lightning as pl
 from omegaconf import DictConfig, OmegaConf
 from pie_datasets import DatasetDict
+from pie_modules.models import *  # noqa: F403
+from pie_modules.taskmodules import *  # noqa: F403
 from pytorch_ie import Pipeline
 
 from src import utils