diff --git a/pygom/model/_ode_composition.py b/pygom/model/_ode_composition.py
index 9f31d13e..5596cb65 100644
--- a/pygom/model/_ode_composition.py
+++ b/pygom/model/_ode_composition.py
@@ -9,6 +9,7 @@
 from functools import reduce
 
 import sympy
+from sympy.matrices import MatrixBase
 import numpy as np
 
 from .base_ode_model import BaseOdeModel
@@ -115,7 +116,7 @@ def generateDirectedDependencyGraph(ode_matrix, transition=None):
         -1 and 1 to indicate the direction of the transition and the state.
         All column sum to one, i.e. transition must have a source and target.
     """
-    assert isinstance(ode_matrix, sympy.matrices.MatrixBase), \
+    assert isinstance(ode_matrix, MatrixBase), \
         "Expecting a vector of expressions"
 
     if transition is None:
@@ -150,7 +151,7 @@ def getUnmatchedExpressionVector(expr_vec, full_output=False):
     list:
         of unmatched expressions, i.e. birth or death processes
     """
-    assert isinstance(expr_vec, sympy.matrices.MatrixBase), \
+    assert isinstance(expr_vec, MatrixBase), \
         "Expecting a vector of expressions"
 
     transition = reduce(lambda x, y: x + y, map(getExpressions, expr_vec))
@@ -181,7 +182,7 @@ def getMatchingExpressionVector(expr_vec, outTuple=False):
     list:
         of matched expressions, i.e. transitions
     """
-    assert isinstance(expr_vec, sympy.matrices.MatrixBase), \
+    assert isinstance(expr_vec, MatrixBase), \
         "Expecting a vector of expressions"
 
     transition = list()
diff --git a/pygom/model/simulate.py b/pygom/model/simulate.py
index d1e8b16a..0abf30f7 100644
--- a/pygom/model/simulate.py
+++ b/pygom/model/simulate.py
@@ -14,6 +14,7 @@
 
 import numpy as np
 import sympy
+import sympy.matrices.matrices
 import scipy.stats
 
 from .deterministic import DeterministicOde