diff --git a/owlapy/owl_neural_reasoners/neural_ontology_manager.py b/owlapy/owl_neural_reasoners/neural_ontology_manager.py
index ab7b85c..7c37c4e 100644
--- a/owlapy/owl_neural_reasoners/neural_ontology_manager.py
+++ b/owlapy/owl_neural_reasoners/neural_ontology_manager.py
@@ -1,4 +1,4 @@
-from owlapy.owl_ontology_manager import OntologyManager
+from owlapy.owl_ontology_manager import OntologyManager, Ontology
 from dicee.knowledge_graph_embeddings import KGE
 
 # LF: for now functions same as normal OntologyManager (+ can return path)
@@ -6,11 +6,11 @@ class NeuralOntologyManager(OntologyManager):
 	def __init__(self, world_store=None):
 		super().__init__(world_store) 
 	
-	def load_ontology(self, path: str = None):
+	def load_ontology(self, path: str = None) -> Ontology:
 		self._path = path
 		return super().load_ontology(path)
 
-	def load_neural_embedding(self, path: str = None):
+	def load_neural_embedding(self, path: str = None) -> KGE:
 		self._path_neural_embedding = path	
 		return KGE(path)
 
diff --git a/owlapy/owl_neural_reasoners/owl_neural_reasoner.py b/owlapy/owl_neural_reasoners/owl_neural_reasoner.py
index 18f1906..97fef3d 100644
--- a/owlapy/owl_neural_reasoners/owl_neural_reasoner.py
+++ b/owlapy/owl_neural_reasoners/owl_neural_reasoner.py
@@ -14,7 +14,6 @@
 from collections import Counter, OrderedDict
 from functools import lru_cache
 from owlapy.owl_neural_reasoners.abstract import AbstractNeuralReasoner
-from owlapy.owl_ontology import Ontology
 from owlapy.owl_neural_reasoners.neural_ontology_manager import NeuralOntologyManager
 
 # TODO: