diff --git a/.github/workflows/lint_test.yml b/.github/workflows/lint_test.yml
index 94cd459..3ada8c6 100644
--- a/.github/workflows/lint_test.yml
+++ b/.github/workflows/lint_test.yml
@@ -11,7 +11,8 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.10' #don't cache this, it will create a cache without all the good stuff
+          python-version: '3.10'
+          cache: 'pip'
       - run: pip install pre-commit pyproject-flake8
       - run: pre-commit install
       - run: pre-commit run --all-files
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 74b4894..e39202e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -27,3 +27,4 @@ repos:
     hooks:
       - id: mypy
         exclude: examples|test|ci
+        additional_dependencies: [torch==2.5.0]  # The version in pyproject.toml must match
diff --git a/emu_base/math/krylov_exp.py b/emu_base/math/krylov_exp.py
index 936a9f6..f14e316 100644
--- a/emu_base/math/krylov_exp.py
+++ b/emu_base/math/krylov_exp.py
@@ -21,7 +21,7 @@ def __init__(
 
 
 def krylov_exp_impl(
-    op: Callable,
+    op: Callable[[torch.Tensor], torch.Tensor],
     v: torch.Tensor,
     is_hermitian: bool,  # note: complex-proportional to its adjoint is enough
     exp_tolerance: float,
@@ -103,7 +103,7 @@ def krylov_exp_impl(
 
 
 def krylov_exp(
-    op: torch.Tensor,
+    op: Callable[[torch.Tensor], torch.Tensor],
     v: torch.Tensor,
     exp_tolerance: float,
     norm_tolerance: float,
diff --git a/emu_base/utils.py b/emu_base/utils.py
index ac36bc5..4a2cba7 100644
--- a/emu_base/utils.py
+++ b/emu_base/utils.py
@@ -1,9 +1,9 @@
 import torch
 
 
-def dist2(left: torch.tensor, right: torch.tensor) -> torch.Tensor:
-    return torch.norm(left - right) ** 2
+def dist2(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+    return torch.pow(torch.norm(left - right), 2)
 
 
-def dist3(left: torch.tensor, right: torch.tensor) -> torch.Tensor:
-    return torch.norm(left - right) ** 3
+def dist3(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+    return torch.pow(torch.norm(left - right), 3)
diff --git a/emu_mps/algebra.py b/emu_mps/algebra.py
index b4ce236..4924c4e 100644
--- a/emu_mps/algebra.py
+++ b/emu_mps/algebra.py
@@ -10,8 +10,8 @@
 
 
 def add_factors(
-    left: list[torch.tensor], right: list[torch.tensor]
-) -> list[torch.tensor]:
+    left: list[torch.Tensor], right: list[torch.Tensor]
+) -> list[torch.Tensor]:
     """
     Direct sum algorithm implementation to sum two tensor trains (MPS/MPO).
     It assumes the left and right bond are along the dimension 0 and -1 of each tensor.
@@ -52,8 +52,8 @@ def add_factors(
 
 
 def scale_factors(
-    factors: list[torch.tensor], scalar: complex, *, which: int
-) -> list[torch.tensor]:
+    factors: list[torch.Tensor], scalar: complex, *, which: int
+) -> list[torch.Tensor]:
     """
     Returns a new list of factors where the tensor at the given index is scaled by `scalar`.
     """
@@ -61,10 +61,10 @@ def scale_factors(
 
 
 def zip_right_step(
-    slider: torch.tensor,
-    top: torch.tensor,
-    bottom: torch.tensor,
-) -> torch.tensor:
+    slider: torch.Tensor,
+    top: torch.Tensor,
+    bottom: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Returns a new `MPS/O` factor of the result of the multiplication MPO @ MPS/O,
     and the updated slider, performing a single step of the
@@ -117,10 +117,10 @@ def zip_right_step(
 
 
 def zip_right(
-    top_factors: list[torch.tensor],
-    bottom_factors: list[torch.tensor],
+    top_factors: list[torch.Tensor],
+    bottom_factors: list[torch.Tensor],
     config: Optional[MPSConfig] = None,
-) -> list[torch.tensor]:
+) -> list[torch.Tensor]:
     """
     Returns a new matrix product, resulting from applying `top` to `bottom`.
     The resulting factors are:
diff --git a/emu_mps/hamiltonian.py b/emu_mps/hamiltonian.py
index c9b4160..d04d18c 100644
--- a/emu_mps/hamiltonian.py
+++ b/emu_mps/hamiltonian.py
@@ -363,7 +363,7 @@ def make_H(
 
     interactions_to_keep = _get_interactions_to_keep(interaction_matrix)
 
-    cores = [_first_factor(interactions_to_keep[0].item())]
+    cores = [_first_factor(interactions_to_keep[0].item() != 0.0)]
 
     if nqubits > 2:
         for i in range(1, middle):
@@ -395,7 +395,7 @@ def make_H(
                 )
             )
     if nqubits == 2:
-        scale = interaction_matrix[0, 1]
+        scale = interaction_matrix[0, 1].item()
     elif interactions_to_keep[-1][0]:
         scale = 1.0
     else:
diff --git a/emu_mps/mpo.py b/emu_mps/mpo.py
index 15163b4..53256aa 100644
--- a/emu_mps/mpo.py
+++ b/emu_mps/mpo.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import itertools
-from typing import Any, List, cast, Iterable, Optional
+from typing import Any, List, Iterable, Optional
 
 import torch
 
@@ -175,14 +175,15 @@ def from_operator_string(
         Returns:
             the operator in MPO form.
         """
+        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
 
         _validate_operator_targets(operations, nqubits)
 
         basis = set(basis)
         if basis == {"r", "g"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "gg": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -197,9 +198,9 @@ def from_operator_string(
                 ).reshape(1, 2, 2, 1),
             }
         elif basis == {"0", "1"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "00": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -218,26 +219,28 @@ def from_operator_string(
 
         mpos = []
         for coeff, tensorop in operations:
-            # this function will recurse through the operators, and replace any definitions
+            # this function will recurse through the operators_with_tensors,
+            # and replace any definitions
             # in terms of strings by the computed tensor
             def replace_operator_string(op: QuditOp | torch.Tensor) -> torch.Tensor:
-                if isinstance(op, dict):
-                    for opstr, coeff in op.items():
-                        tensor = replace_operator_string(operators[opstr])
-                        operators[opstr] = tensor
-                        op[opstr] = tensor * coeff
-                    op = sum(cast(list[torch.Tensor], op.values()))
-                return op
+                if isinstance(op, torch.Tensor):
+                    return op
+
+                result = torch.zeros(1, 2, 2, 1, dtype=torch.complex128)
+                for opstr, coeff in op.items():
+                    tensor = replace_operator_string(operators_with_tensors[opstr])
+                    operators_with_tensors[opstr] = tensor
+                    result += tensor * coeff
+                return result
 
             factors = [
                 torch.eye(2, 2, dtype=torch.complex128).reshape(1, 2, 2, 1)
             ] * nqubits
 
-            for i, op in enumerate(tensorop):
-                tensorop[i] = (replace_operator_string(op[0]), op[1])
-
             for op in tensorop:
-                for i in op[1]:
-                    factors[i] = op[0]
+                factor = replace_operator_string(op[0])
+                for target_qubit in op[1]:
+                    factors[target_qubit] = factor
+
             mpos.append(coeff * MPO(factors, **kwargs))
         return sum(mpos[1:], start=mpos[0])
diff --git a/emu_mps/mps.py b/emu_mps/mps.py
index 82e97df..84b798d 100644
--- a/emu_mps/mps.py
+++ b/emu_mps/mps.py
@@ -226,7 +226,7 @@ def _sample_implementation(self, rnd_vector: torch.Tensor) -> str:
         num_qubits = len(self.factors)
 
         bitstring = ""
-        acc_mps_j: torch.tensor = self.factors[0]
+        acc_mps_j: torch.Tensor = self.factors[0]
 
         for qubit in range(num_qubits):
             # comp_basis is a projector: 0 is for ket |0> and 1 for ket |1>
diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py
index 595f8a0..6c50d58 100644
--- a/emu_mps/mps_backend_impl.py
+++ b/emu_mps/mps_backend_impl.py
@@ -298,7 +298,7 @@ def progress(self) -> None:
             )
             if not self.has_lindblad_noise:
                 # Free memory because it won't be used anymore
-                self.right_baths[-2] = None
+                self.right_baths[-2] = torch.zeros(0)
 
             self._evolve(self.tdvp_index, dt=-delta_time / 2)
             self.left_baths.pop()
@@ -453,13 +453,12 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
     """
 
     jump_threshold: float
-    aggregated_lindblad_ops: Optional[torch.Tensor]
+    aggregated_lindblad_ops: torch.Tensor
     norm_gap_before_jump: float
     root_finder: Optional[BrentsRootFinder]
 
     def __init__(self, config: MPSConfig, pulser_data: PulserData):
         super().__init__(config, pulser_data)
-        self.aggregated_lindblad_ops = None
         self.lindblad_ops = pulser_data.lindblad_ops
         self.root_finder = None
 
@@ -520,7 +519,7 @@ def do_random_quantum_jump(self) -> None:
                 for qubit in range(self.state.num_sites)
                 for op in self.lindblad_ops
             ],
-            weights=jump_operator_weights.reshape(-1),
+            weights=jump_operator_weights.reshape(-1).tolist(),
         )[0]
 
         self.state.apply(jumped_qubit_index, jump_operator)
diff --git a/emu_mps/tdvp.py b/emu_mps/tdvp.py
index d80e70b..d4dd429 100644
--- a/emu_mps/tdvp.py
+++ b/emu_mps/tdvp.py
@@ -11,7 +11,8 @@ def new_right_bath(
 ) -> torch.Tensor:
     bath = torch.tensordot(state, bath, ([2], [2]))
     bath = torch.tensordot(op.to(bath.device), bath, ([2, 3], [1, 3]))
-    return torch.tensordot(state.conj(), bath, ([1, 2], [1, 3]))
+    bath = torch.tensordot(state.conj(), bath, ([1, 2], [1, 3]))
+    return bath
 
 
 """
diff --git a/emu_mps/utils.py b/emu_mps/utils.py
index 5843dee..3add64f 100644
--- a/emu_mps/utils.py
+++ b/emu_mps/utils.py
@@ -12,15 +12,16 @@ def new_left_bath(
     # this order is more efficient than contracting the op first in general
     bath = torch.tensordot(bath, state.conj(), ([0], [0]))
     bath = torch.tensordot(bath, op.to(bath.device), ([0, 2], [0, 1]))
-    return torch.tensordot(bath, state, ([0, 2], [0, 1]))
+    bath = torch.tensordot(bath, state, ([0, 2], [0, 1]))
+    return bath
 
 
 def _determine_cutoff_index(d: torch.Tensor, max_error: float) -> int:
     assert max_error > 0
     squared_max_error = max_error * max_error
-    acc = 0
+    acc = 0.0
     for i in range(d.shape[0]):
-        acc += d[i]
+        acc += d[i].item()
         if acc > squared_max_error:
             return i
     return 0  # type: ignore[no-any-return]
@@ -60,7 +61,7 @@ def split_tensor(
 
 
 def truncate_impl(
-    factors: list[torch.tensor],
+    factors: list[torch.Tensor],
     config: MPSConfig,
 ) -> None:
     """
diff --git a/emu_sv/dense_operator.py b/emu_sv/dense_operator.py
index ce8ba46..36f7db8 100644
--- a/emu_sv/dense_operator.py
+++ b/emu_sv/dense_operator.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import itertools
-from typing import Any, cast, Iterable
+from typing import Any, Iterable
 
 import torch
 from emu_base.base_classes.operator import FullOp, QuditOp
@@ -114,7 +114,7 @@ def expect(self, state: State) -> float | complex:
         ), "currently, only expectation values of StateVectors are \
         supported"
 
-        return torch.vdot(state.vector, self.matrix @ state.vector)  # type: ignore [no-any-return]
+        return torch.vdot(state.vector, self.matrix @ state.vector).item()
 
     @staticmethod
     def from_operator_string(
@@ -140,20 +140,22 @@ def from_operator_string(
 
         _validate_operator_targets(operations, nqubits)
 
+        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
+
         basis = set(basis)
         if basis == {"r", "g"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "gg": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128),
                 "gr": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128),
                 "rg": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128),
                 "rr": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.complex128),
             }
         elif basis == {"0", "1"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "00": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128),
                 "01": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128),
                 "10": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128),
@@ -164,29 +166,29 @@ def from_operator_string(
 
         accum_res = torch.zeros(2**nqubits, 2**nqubits, dtype=torch.complex128)
         for coeff, tensorop in operations:
-            # this function will recurse through the operators, and replace any definitions
-            # in terms of strings by the computed matrix
+            # this function will recurse through the operators_with_tensors,
+            # and replace any definitions in terms of strings by the computed matrix
             def replace_operator_string(op: QuditOp | torch.Tensor) -> torch.Tensor:
-                if isinstance(op, dict):
-                    for opstr, coeff in op.items():
-                        tensor = replace_operator_string(operators[opstr])
-                        operators[opstr] = tensor
-                        op[opstr] = tensor * coeff
-                    op = sum(cast(list[torch.Tensor], op.values()))
-                return op
+                if isinstance(op, torch.Tensor):
+                    return op
 
-            pre_dense_op = [torch.eye(2, 2, dtype=torch.complex128)] * nqubits
+                result = torch.zeros(2, 2, dtype=torch.complex128)
+                for opstr, coeff in op.items():
+                    tensor = replace_operator_string(operators_with_tensors[opstr])
+                    operators_with_tensors[opstr] = tensor
+                    result += tensor * coeff
+                return result
 
-            for i, op in enumerate(tensorop):
-                tensorop[i] = (replace_operator_string(op[0]), op[1])
+            total_op_per_qubit = [torch.eye(2, 2, dtype=torch.complex128)] * nqubits
 
             for op in tensorop:
-                for i in op[1]:
-                    pre_dense_op[i] = op[0]
+                factor = replace_operator_string(op[0])
+                for target_qubit in op[1]:
+                    total_op_per_qubit[target_qubit] = factor
 
-            dense_op = pre_dense_op[0]
-            for i in pre_dense_op[1:]:
-                dense_op = torch.kron(dense_op, i)
+            dense_op = total_op_per_qubit[0]
+            for single_qubit_operator in total_op_per_qubit[1:]:
+                dense_op = torch.kron(dense_op, single_qubit_operator)
 
             accum_res += coeff * dense_op
         return DenseOperator(accum_res)
diff --git a/emu_sv/hamiltonian.py b/emu_sv/hamiltonian.py
index 65b01ab..09e2398 100644
--- a/emu_sv/hamiltonian.py
+++ b/emu_sv/hamiltonian.py
@@ -47,7 +47,7 @@ def __init__(
         omegas: torch.Tensor,
         deltas: torch.Tensor,
         interaction_matrix: torch.Tensor,
-        device: torch.DeviceObjType,
+        device: torch.device,
     ):
         self.nqubits: int = len(omegas)
         self.omegas: torch.Tensor = omegas / 2.0
diff --git a/emu_sv/state_vector.py b/emu_sv/state_vector.py
index 5b9e8b4..105ceb8 100644
--- a/emu_sv/state_vector.py
+++ b/emu_sv/state_vector.py
@@ -25,12 +25,15 @@ class StateVector(State):
         vector (torch.Tensor): 1D tensor representation of a state vector.
 
     Methods:
-        __init__(vector: torch.Tensor):
+        __init__(vector: torch.Tensor, gpu: bool = False):
             Initializes the state vector. Ensures the length is a power of 2.
 
-        _normalize() -> torch.Tensor:
+        _normalize() -> None:
             Normalizes the state vector to ensure it represents a valid quantum state.
 
+        zero(num_sites: int, gpu: bool = False) -> StateVector:
+            Creates a zero state vector for a specified number of qubits.
+
         make(num_sites: int, gpu: bool = False) -> StateVector:
             Creates a ground state vector |000...0> for a specified number of qubits.
 
@@ -75,6 +78,8 @@ class StateVector(State):
     def __init__(
         self,
         vector: torch.Tensor,
+        *,
+        gpu: bool = False,
     ):
         # NOTE: this accepts also zero vectors.
 
@@ -82,16 +87,36 @@ def __init__(
             len(vector)
         ).is_integer(), "The number of elements in the vector should be power of 2"
 
-        self.vector = vector.to(dtype=dtype)
+        device = "cuda" if gpu else "cpu"
+        self.vector = vector.to(dtype=dtype, device=device)
 
-    def _normalize(self) -> torch.Tensor:
+    def _normalize(self) -> None:
         # NOTE: use this in the callbacks
         """Checks if the input is normalized or not"""
         norm_state = torch.linalg.vector_norm(self.vector)
 
         if not torch.allclose(norm_state, torch.tensor(1.0, dtype=torch.float64)):
             self.vector = self.vector / norm_state
-        return self.vector
+
+    @classmethod
+    def zero(cls, num_sites: int, gpu: bool = False) -> StateVector:
+        """
+        Returns a zero State vector.
+        The vector in the output of StateVector has the shape (2,)*number of qubits
+
+        Args:
+            num_sites: the number of qubits
+            gpu: whether gpu or cpu
+
+        Example:
+        -------
+        >>> StateVector.zero(2)
+        tensor([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=torch.complex128)
+        """
+
+        device = "cuda" if gpu else "cpu"
+        vector = torch.zeros(2**num_sites, dtype=dtype, device=device)
+        return cls(vector, gpu=gpu)
 
     @classmethod
     def make(cls, num_sites: int, gpu: bool = False) -> StateVector:
@@ -101,7 +126,7 @@ def make(cls, num_sites: int, gpu: bool = False) -> StateVector:
 
         Args:
             num_sites: the number of qubits
-            gpu: weather gpu or cpu
+            gpu: whether gpu or cpu
 
         Example:
         -------
@@ -109,12 +134,11 @@ def make(cls, num_sites: int, gpu: bool = False) -> StateVector:
         tensor([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=torch.complex128)
         """
 
-        device = "cuda" if gpu else "cpu"
-        ground_state = torch.zeros(2**num_sites, dtype=dtype, device=device)
-        ground_state[0] = torch.tensor(1.0, dtype=dtype, device=device)
-        return cls(ground_state)
+        result = cls.zero(num_sites=num_sites, gpu=gpu)
+        result.vector[0] = 1.0
+        return result
 
-    def inner(self, other: State) -> torch.Tensor:
+    def inner(self, other: State) -> float | complex:
         assert isinstance(
             other, StateVector
         ), "Other state also needs to be a StateVector"
@@ -122,14 +146,14 @@ def inner(self, other: State) -> torch.Tensor:
             self.vector.shape == other.vector.shape
         ), "States do not have the same number of sites"
 
-        return torch.vdot(self.vector, other.vector)
+        return torch.vdot(self.vector, other.vector).item()
 
     def sample(
         self, num_shots: int = 1000, p_false_pos: float = 0.0, p_false_neg: float = 0.0
     ) -> Counter[str]:
         """Probability distribution over measurement outcomes"""
 
-        probabilities = torch.abs((self.vector)) ** 2
+        probabilities = torch.abs(self.vector) ** 2
         probabilities /= probabilities.sum()  # multinomial does not normalize the input
 
         outcomes = torch.multinomial(probabilities, num_shots, replacement=True)
@@ -161,12 +185,12 @@ def __rmul__(self, scalar: complex) -> StateVector:
 
         return StateVector(result)
 
-    def norm(self) -> torch.tensor:
+    def norm(self) -> float | complex:
         """Norm of the state"""
-        return torch.linalg.norm(self.vector)
+        norm: float | complex = torch.linalg.vector_norm(self.vector).item()
+        return norm
 
     def __repr__(self) -> str:
-
         return repr(self.vector)
 
     @staticmethod
@@ -175,7 +199,6 @@ def from_state_string(
         basis: Iterable[str],
         nqubits: int,
         strings: dict[str, complex],
-        gpu: bool = False,
         **kwargs: Any,
     ) -> StateVector:
         """Transforms a state given by a string into a state vector.
@@ -187,7 +210,6 @@ def from_state_string(
             basis: A tuple containing the basis states (e.g., ('r', 'g')).
             nqubits: the number of qubits.
             strings: A dictionary mapping state strings to complex or floats amplitudes.
-            gpu: weather gpu or cpu. By deafult, cpu
 
         Returns:
             The resulting state.
@@ -202,8 +224,6 @@ def from_state_string(
                dtype=torch.complex128)
         """
 
-        device = "gpu" if gpu else "cpu"
-
         basis = set(basis)
         if basis == {"r", "g"}:
             one = "r"
@@ -212,14 +232,13 @@ def from_state_string(
         else:
             raise ValueError("Unsupported basis provided")
 
-        accum_zero = torch.zeros(2**nqubits, dtype=dtype, device=device)
-        accum_state = StateVector(accum_zero)
+        accum_state = StateVector.zero(num_sites=nqubits, **kwargs)
 
         for state, amplitude in strings.items():
             bin_to_int = int(
                 state.replace(one, "1").replace("g", "0"), 2
             )  # "0" basis is already in "0"
-            accum_state.vector[bin_to_int] = amplitude
+            accum_state.vector[bin_to_int] = torch.tensor([amplitude])
 
         accum_state._normalize()
 
diff --git a/pyproject.toml b/pyproject.toml
index 94907d9..35a89bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ classifiers=[
 ]
 dependencies = [
   "pulser-core==1.2.*",
-  "torch==2.5.0"]
+  "torch==2.5.0"]  # The version in .pre-commit-config.yaml must match
 dynamic = ["version"]
 
 [tool.setuptools]
diff --git a/test_dependency_versions.sh b/test_dependency_versions.sh
index 2fb36be..6e9d2ca 100755
--- a/test_dependency_versions.sh
+++ b/test_dependency_versions.sh
@@ -1,18 +1,32 @@
 #!/bin/bash
+set -e
 
-mps_string=`grep emu-base ci/emu_mps/pyproject.toml`
+mps_string="$(grep emu-base ci/emu_mps/pyproject.toml)"
 [[ "$mps_string" =~ .*([0-9]\.[0-9]\.[0-9]).* ]]
-mps_dep=${BASH_REMATCH[1]}
-echo emu-mps depends on emu-base version $mps_dep
+mps_dep="${BASH_REMATCH[1]}"
+echo "emu-mps depends on emu-base version $mps_dep"
 
-base_string=`grep version emu_base/__init__.py`
+base_string="$(grep version emu_base/__init__.py)"
 [[ "$base_string" =~ .*([0-9]\.[0-9]\.[0-9]).* ]]
-base_version=${BASH_REMATCH[1]}
-echo emu-base is version $base_version
+base_version="${BASH_REMATCH[1]}"
+echo "emu-base is version $base_version"
 
-if [[ $mps_dep == $base_version ]]
+if [[ "$mps_dep" != "$base_version" ]]
 then
-exit 0
-else
-exit 1
+    exit 1
+fi
+
+torch_pre_commit_string="$(grep torch .pre-commit-config.yaml)"
+[[ "$torch_pre_commit_string" =~ .*([0-9]\.[0-9]\.[0-9]).* ]]
+torch_pre_commit_version="${BASH_REMATCH[1]}"
+echo "pre-commit uses torch version $torch_pre_commit_version"
+
+torch_pyproject_string="$(grep torch pyproject.toml)"
+[[ "$torch_pyproject_string" =~ .*([0-9]\.[0-9]\.[0-9]).* ]]
+torch_pyproject_version="${BASH_REMATCH[1]}"
+echo "pyproject.toml uses torch version $torch_pyproject_version"
+
+if [[ "$torch_pre_commit_version" != "$torch_pyproject_version" ]]
+then
+    exit 1
 fi