diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py
index c0dac32..527cb38 100644
--- a/src/src_method/_sweep.py
+++ b/src/src_method/_sweep.py
@@ -19,7 +19,7 @@
 
 import logging
 from time import perf_counter_ns
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -64,8 +64,9 @@ def sweep(
     *,
     cutoff: float = 0.0,
     dtype: DTypeLike,
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Contract and compress a stack in ket form with one SRC sweep.
 
     Args:
@@ -80,11 +81,13 @@ def sweep(
         xp: Array module (``numpy`` or ``cupy``).
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
         dtype: The data type of the sketches.
+        to_host: If ``True`` (default), the result is copied to the host as numpy
+            arrays. If ``False``, it stays on the device of ``xp``.
         resources: Memory budgets and scratch space; detected when ``None``.
 
     Returns:
         The site arrays of the compressed train in right-canonical form, as numpy
-        arrays.
+        arrays if ``to_host`` is ``True``, otherwise as arrays of ``xp``.
     """
     shapes = padded_shapes(layers)
     # The dtype of the environments and the output: the sketches promoted by the
@@ -140,7 +143,12 @@ def sweep(
             store.stall_seconds,
         )
         logger.debug("Device pool: %d B", device_pool_bytes(xp))
-    return [to_numpy(site) for site in unpad(eta, kind)]
+
+    # La tua fix per il GPU round-trip applicata al nuovo codice:
+    sites_out = unpad(eta, kind)
+    if to_host:
+        return [to_numpy(site) for site in sites_out]
+    return list(sites_out)
 
 
 def _left_to_right(
@@ -211,10 +219,10 @@ def _right_to_left(
         for lo, hi in _ranges(Q.shape[1], site.project_batch):
             new_proj[lo:hi] = kernels.project(eta_j[lo:hi].conj(), cores, proj)
         proj = new_proj
-        eta_reversed.append(to_numpy(eta_j))
+        eta_reversed.append(eta_j)
     cores = source[0]
     up, down = cores[0].shape[2], cores[-1].shape[3]
     first = xp.empty((proj.shape[0], up, down), dtype=dtype)
     for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch):
         first[lo:hi] = kernels.first(cores, proj[lo:hi]).reshape(hi - lo, up, down)
-    return [to_numpy(first)[None], *reversed(eta_reversed)]
+    return [first[None], *reversed(eta_reversed)]
diff --git a/src/src_method/_tensor_train.py b/src/src_method/_tensor_train.py
index 5e69020..5ad2de6 100644
--- a/src/src_method/_tensor_train.py
+++ b/src/src_method/_tensor_train.py
@@ -22,10 +22,11 @@
 from opt_einsum import contract
 
 from ._validation import validate_open_boundary
-from .utils import to_numpy
+from .utils import is_host, to_numpy
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from types import ModuleType
 
     from numpy.typing import NDArray
 
@@ -153,8 +154,18 @@ def check_exact_supported(n_sites: int) -> None:
         raise ValueError(msg)
 
 
+def _to_backend(site: Site, xp: ModuleType) -> NDArray:
+    """Ensure a site tensor is an array on the backend ``xp``."""
+    if is_host(xp):
+        return to_numpy(site)
+    return site if isinstance(site, xp.ndarray) else xp.asarray(site)
+
+
 def exact_compress(
-    arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
+    arrays: Sequence[Site],
+    chi_out: int,
+    kind: TrainKind,
+    xp: ModuleType = np,
 ) -> list[NDArray]:
     """Compress a two-site train exactly via a single truncated SVD.
 
@@ -164,22 +175,24 @@ def exact_compress(
         arrays: The two site tensors of the train.
         chi_out: The maximum bond dimension to keep.
         kind: Whether the train is an ``"mps"`` or an ``"mpo"``.
+        xp: Array module (``numpy`` or ``cupy``); defaults to numpy.
 
     Returns:
-        The compressed train, in right-canonical form, as numpy arrays.
+        The compressed train, in right-canonical form, as arrays of ``xp``.
     """
-    # The dense SVD is host-side, so accept device arrays like the sweep does.
-    arrays = [to_numpy(arr) for arr in arrays]
+    device_arrays = [_to_backend(arr, xp) for arr in arrays]
     if kind == "mps":
         # (b, p0) x (b, p1) -> (p0, p1)
-        theta = contract("ab,ac->bc", arrays[0], arrays[1])
-        left, right = _truncated_svd(theta, chi_out)
+        theta = contract("ab,ac->bc", device_arrays[0], device_arrays[1])
+        left, right = _truncated_svd(theta, chi_out, xp=xp)
         return [left.T, right]
 
     # (b, u0, d0) x (b, u1, d1) -> (u0, d0, u1, d1)
-    theta = contract("aij,akl->ijkl", arrays[0], arrays[1])
+    theta = contract("aij,akl->ijkl", device_arrays[0], device_arrays[1])
     up_l, down_l, up_r, down_r = theta.shape
-    left, right = _truncated_svd(theta.reshape(up_l * down_l, up_r * down_r), chi_out)
+    left, right = _truncated_svd(
+        theta.reshape(up_l * down_l, up_r * down_r), chi_out, xp=xp
+    )
     rank = left.shape[1]
     return [
         left.reshape(up_l, down_l, rank).transpose(2, 0, 1),
@@ -187,9 +200,11 @@ def exact_compress(
     ]
 
 
-def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]:
+def _truncated_svd(
+    theta: NDArray, chi_out: int, xp: ModuleType = np
+) -> tuple[NDArray, NDArray]:
     """Split a matrix as ``(U @ diag(S), Vh)``, keeping at most ``chi_out`` values."""
-    U, S, Vh = np.linalg.svd(theta, full_matrices=False)
+    U, S, Vh = xp.linalg.svd(theta, full_matrices=False)
     rank = min(chi_out, S.size)
     return U[:, :rank] * S[:rank], Vh[:rank]
 
@@ -397,7 +412,10 @@ def _check_physical_dims(
 
 
 def exact_stack(
-    layers: Sequence[Sequence[Site]], chi_out: int, kind: TrainKind
+    layers: Sequence[Sequence[Site]],
+    chi_out: int,
+    kind: TrainKind,
+    xp: ModuleType = np,
 ) -> list[NDArray]:
     """Contract and compress a two-site stack exactly.
 
@@ -409,15 +427,16 @@ def exact_stack(
         layers: A stack in ket form, as returned by `normalize_stack`.
         chi_out: The maximum bond dimension to keep.
         kind: The kind of the contracted train.
+        xp: Array module (``numpy`` or ``cupy``); defaults to numpy.
 
     Returns:
-        The compressed product, in right-canonical form, as numpy arrays.
+        The compressed product, in right-canonical form, as arrays of ``xp``.
     """
-    padded = [pad([to_numpy(site) for site in layer]) for layer in layers]
+    padded = [pad([_to_backend(site, xp) for site in layer]) for layer in layers]
     product = padded[-1]
     for layer in reversed(padded[:-1]):
         product = [_fuse(upper, lower) for upper, lower in zip(layer, product)]
-    return exact_compress(unpad(product, kind), chi_out, kind)
+    return exact_compress(unpad(product, kind), chi_out, kind, xp=xp)
 
 
 def _fuse(upper: NDArray, lower: NDArray) -> NDArray:
diff --git a/src/src_method/apply.py b/src/src_method/apply.py
index 1fd6969..6bd3a44 100644
--- a/src/src_method/apply.py
+++ b/src/src_method/apply.py
@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from ._tensor_train import infer_kind
 from .stack import src
@@ -17,7 +17,7 @@
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -32,8 +32,9 @@ def apply(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
     Equivalent to ``src(left_tensor, right_tensor, ...)`` restricted to an MPO on
@@ -59,6 +60,8 @@ def apply(
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy).  Requires
             the optional ``cupy`` dependency for GPU execution.
+        to_host: If ``True`` (default), results are returned as numpy arrays on
+            the host. Set to ``False`` to keep tensors on the compute device.
         resources: Memory budgets and scratch space; see `src_method.stack.src`.
 
     Returns:
@@ -91,5 +94,6 @@ def apply(
         dtype=dtype,
         seed=seed,
         device=device,
+        to_host=to_host,
         resources=resources,
     )
diff --git a/src/src_method/compress.py b/src/src_method/compress.py
index 46e7ffe..fff5f69 100644
--- a/src/src_method/compress.py
+++ b/src/src_method/compress.py
@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from .stack import src
 
@@ -16,7 +16,7 @@
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -30,8 +30,9 @@ def compress(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
     Equivalent to ``src(tensor, ...)``; see `src_method.stack.src` for the
@@ -55,6 +56,8 @@ def compress(
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy).  Requires
             the optional ``cupy`` dependency for GPU execution.
+        to_host: If ``True`` (default), results are returned as numpy arrays on
+            the host. Set to ``False`` to keep tensors on the compute device.
         resources: Memory budgets and scratch space; see `src_method.stack.src`.
 
     Returns:
@@ -76,5 +79,6 @@ def compress(
         dtype=dtype,
         seed=seed,
         device=device,
+        to_host=to_host,
         resources=resources,
     )
diff --git a/src/src_method/stack.py b/src/src_method/stack.py
index 631e314..5c05881 100644
--- a/src/src_method/stack.py
+++ b/src/src_method/stack.py
@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from ._sweep import sweep
 from ._tensor_train import (
@@ -18,13 +18,13 @@
     normalize_stack,
 )
 from ._validation import validate_chi_out, validate_cutoff
-from .utils import default_rng, get_xp, sketch_dtype
+from .utils import default_rng, get_xp, sketch_dtype, to_numpy
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -44,8 +44,9 @@ def src(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Contract a stack of tensor trains and compress the result with SRC.
 
     The stack ``src(T_1, T_2, ..., T_m)`` is the product ``T_1 T_2 ... T_m`` in
@@ -88,13 +89,15 @@ def src(
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires
             the optional ``cupy`` dependency for GPU execution.
+        to_host: If ``True`` (default), results are returned as numpy arrays on
+            the host. Set to ``False`` to keep tensors on the compute device.
         resources: Memory budgets and scratch space for the sweep (see
             `Resources`); every budget left unset is detected. Ignored for
             two-site stacks.
 
     Returns:
         The site arrays of the compressed train (MPS or MPO), in right-canonical
-        form, as numpy arrays (host-side, whatever the ``device``).
+        form, as numpy arrays if ``to_host`` is ``True``, otherwise as arrays of ``xp``.
 
     Raises:
         TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised
@@ -116,7 +119,10 @@ def src(
     if n_sites < MIN_SRC_SITES:
         check_exact_supported(n_sites)
         logger.warning(LOG_WARN_SMALL)
-        return exact_stack(layers, chi_out, kind)
+        result = exact_stack(layers, chi_out, kind, xp=xp)
+        if to_host:
+            return [to_numpy(site) for site in result]
+        return list(result)
 
     logger.debug(
         "Starting SRC: n_sites=%d, depth=%d, output=%s, device=%s",
@@ -133,6 +139,7 @@ def src(
         xp,
         cutoff=cutoff,
         dtype=sketch_dtype(dtype, *layers),
+        to_host=to_host,
         resources=resources,
     )
     logger.debug("SRC complete")
diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py
index ff2e89c..bbc5935 100644
--- a/src/src_method/utils/linalg.py
+++ b/src/src_method/utils/linalg.py
@@ -40,16 +40,13 @@ def truncated_qr(
         return xp.linalg.qr(matrix)[0]
 
     Q, R = xp.linalg.qr(matrix.T if transpose else matrix)
-    R_np = R.get() if hasattr(R, "get") else np.asarray(R)
-    U, S, _ = np.linalg.svd(R_np.T if transpose else R_np, full_matrices=False)
+    U, S, _ = xp.linalg.svd(R.T if transpose else R, full_matrices=False)
     rank = max(1, int((cutoff * S[0] <= S).sum()))
 
     if transpose:
         Q_trunc = U[:, :rank]
     else:
-        U_trunc = xp.asarray(U[:, :rank]) if xp is not np else U[:, :rank]
+        U_trunc = U[:, :rank]
         Q_trunc = Q @ U_trunc
 
-    if xp is not np:
-        return xp.asarray(Q_trunc)
     return Q_trunc
diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py
index 6b4fdb6..37d2f1e 100644
--- a/tests/test_gpu_backend.py
+++ b/tests/test_gpu_backend.py
@@ -22,6 +22,7 @@
     make_plan,
 )
 from src_method._sites import padded_shapes, site_bytes
+from src_method.utils import truncated_qr
 
 cupy = pytest.importorskip("cupy")
 
@@ -157,6 +158,92 @@ def test_compress_mpo_cpu_gpu_equivalent() -> None:
     np.testing.assert_allclose(cpu_out.distance(gpu_out), 0.0, atol=1e-6)
 
 
+def test_apply_mpo_mps_cutoff_cpu_gpu_equivalent() -> None:
+    """CPU and GPU ``apply`` MPO @ MPS with positive cutoff must produce equivalent MPSs."""
+    n_sites, phys_dim, chi_out = 5, 2, 8
+    H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    psi = qtn.MPS_rand_state(
+        n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
+    )
+
+    cpu_out = as_mps(
+        apply(
+            H.arrays,
+            psi.arrays,
+            chi_out=chi_out,
+            cutoff=1e-4,
+            dtype=np.complex128,
+            seed=42,
+            device="cpu",
+        )
+    )
+    gpu_out = as_mps(
+        apply(
+            H.arrays,
+            psi.arrays,
+            chi_out=chi_out,
+            cutoff=1e-4,
+            dtype=np.complex128,
+            seed=42,
+            device="gpu",
+        )
+    )
+
+    np.testing.assert_allclose(cpu_out.distance(gpu_out), 0.0, atol=1e-6)
+
+
+def test_compress_mpo_cutoff_cpu_gpu_equivalent() -> None:
+    """CPU and GPU ``compress`` with positive cutoff must produce equivalent MPOs."""
+    n_sites, phys_dim, chi_out = 8, 4, 32
+    A = qtn.MPO_rand(n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128)
+
+    cpu_out = as_mpo(
+        compress(
+            A.arrays,
+            chi_out=chi_out,
+            cutoff=1e-4,
+            dtype=np.complex128,
+            seed=42,
+            device="cpu",
+        )
+    )
+    gpu_out = as_mpo(
+        compress(
+            A.arrays,
+            chi_out=chi_out,
+            cutoff=1e-4,
+            dtype=np.complex128,
+            seed=42,
+            device="gpu",
+        )
+    )
+
+    np.testing.assert_allclose(cpu_out.distance(gpu_out), 0.0, atol=1e-6)
+
+
+@pytest.mark.parametrize("shape", [(20, 10), (10, 20)])
+def test_truncated_qr_positive_cutoff_gpu_matches_cpu(shape: tuple[int, int]) -> None:
+    """``truncated_qr`` with positive cutoff on GPU must compute on device and match CPU."""
+    rng = np.random.default_rng(42)
+    matrix = rng.normal(size=shape) + 1j * rng.normal(size=shape)
+    gpu_matrix = cupy.asarray(matrix)
+
+    cpu_q = truncated_qr(matrix, cutoff=0.1, xp=np)
+    gpu_q = truncated_qr(gpu_matrix, cutoff=0.1, xp=cupy)
+
+    assert isinstance(gpu_q, cupy.ndarray)
+    assert gpu_q.shape == cpu_q.shape
+
+    # Both must be isometries: Q^H @ Q == I
+    identity = cupy.eye(gpu_q.shape[1], dtype=gpu_q.dtype)
+    cupy.testing.assert_allclose(gpu_q.conj().T @ gpu_q, identity, atol=1e-12)
+
+    # Subspaces must match: compare projectors P = Q @ Q^H
+    proj_cpu = cpu_q @ cpu_q.conj().T
+    proj_gpu = (gpu_q @ gpu_q.conj().T).get()
+    np.testing.assert_allclose(proj_gpu, proj_cpu, atol=1e-12)
+
+
 def test_invalid_device_raises() -> None:
     """An unrecognised ``device`` must raise ``ValueError``."""
     H = qtn.MPO_rand(5, bond_dim=4, phys_dim=2, dtype=np.complex128)
@@ -195,6 +282,189 @@ def test_src_stack_gpu_matches_reference(device: str) -> None:
     np.testing.assert_allclose(ref.distance(out), 0.0, atol=1e-6)
 
 
+def test_apply_mpo_mps_to_host_false_returns_cupy() -> None:
+    """With to_host=False, apply() must return cupy arrays on GPU."""
+    n_sites, phys_dim, chi_out = 5, 2, 8
+    H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    psi = qtn.MPS_rand_state(
+        n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
+    )
+
+    result = apply(
+        H.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="gpu",
+        to_host=False,
+    )
+    assert all(isinstance(t, cupy.ndarray) for t in result)
+
+
+def test_compress_mpo_to_host_false_returns_cupy() -> None:
+    """With to_host=False, compress() must return cupy arrays on GPU."""
+    n_sites, phys_dim, chi_out = 5, 2, 8
+    H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+
+    result = compress(
+        H.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="gpu",
+        to_host=False,
+    )
+    assert all(isinstance(t, cupy.ndarray) for t in result)
+
+
+def test_chained_apply_no_host_roundtrip() -> None:
+    """Chained GPU calls with to_host=False must work and match CPU results."""
+    n_sites, phys_dim, chi_out = 5, 2, 8
+    H1 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    H2 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    psi = qtn.MPS_rand_state(
+        n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
+    )
+
+    # 1. Chained GPU apply
+    intermediate_gpu = apply(
+        H1.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="gpu",
+        to_host=False,
+    )
+    assert isinstance(intermediate_gpu[0], cupy.ndarray)
+
+    final_gpu = apply(
+        H2.arrays,
+        intermediate_gpu,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="gpu",
+        to_host=False,
+    )
+    assert isinstance(final_gpu[0], cupy.ndarray)
+
+    # 2. Chained CPU apply for numerical comparison
+    intermediate_cpu = apply(
+        H1.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="cpu",
+        to_host=True,
+    )
+    final_cpu = apply(
+        H2.arrays,
+        intermediate_cpu,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="cpu",
+        to_host=True,
+    )
+
+    # Bring GPU result to host and compare
+    final_gpu_host = [t.get() for t in final_gpu]
+
+    qtn_cpu = as_mps(final_cpu)
+    qtn_gpu = as_mps(final_gpu_host)
+    np.testing.assert_allclose(qtn_cpu.distance(qtn_gpu), 0.0, atol=1e-6)
+
+
+def test_to_host_true_backward_compatible() -> None:
+    """to_host=True (implicit default) must return numpy arrays even when device='gpu'."""
+    n_sites, phys_dim, chi_out = 5, 2, 8
+    H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+
+    result = compress(
+        H.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        seed=0,
+        device="gpu",
+        # to_host=True is the implicit default here
+    )
+    assert all(isinstance(t, np.ndarray) for t in result)
+
+
+def test_two_site_fallback_to_host_false_returns_cupy() -> None:
+    """Two-site fallback with to_host=False must return cupy arrays on GPU."""
+    n_sites, phys_dim, chi_out = 2, 2, 4
+    H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    psi = qtn.MPS_rand_state(
+        n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
+    )
+
+    result = apply(
+        H.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        device="gpu",
+        to_host=False,
+    )
+    assert all(isinstance(t, cupy.ndarray) for t in result)
+
+
+def test_two_site_chained_apply_no_host_roundtrip() -> None:
+    """Chained two-site GPU calls with to_host=False must match CPU results."""
+    n_sites, phys_dim, chi_out = 2, 2, 4
+    H1 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    H2 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128)
+    psi = qtn.MPS_rand_state(
+        n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
+    )
+
+    intermediate_gpu = apply(
+        H1.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        device="gpu",
+        to_host=False,
+    )
+    assert isinstance(intermediate_gpu[0], cupy.ndarray)
+
+    final_gpu = apply(
+        H2.arrays,
+        intermediate_gpu,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        device="gpu",
+        to_host=False,
+    )
+    assert isinstance(final_gpu[0], cupy.ndarray)
+
+    intermediate_cpu = apply(
+        H1.arrays,
+        psi.arrays,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        device="cpu",
+        to_host=True,
+    )
+    final_cpu = apply(
+        H2.arrays,
+        intermediate_cpu,
+        chi_out=chi_out,
+        dtype=np.complex128,
+        device="cpu",
+        to_host=True,
+    )
+
+    final_gpu_host = [t.get() for t in final_gpu]
+    qtn_cpu = as_mps(final_cpu)
+    qtn_gpu = as_mps(final_gpu_host)
+    np.testing.assert_allclose(qtn_cpu.distance(qtn_gpu), 0.0, atol=1e-6)
+
+
 # ---------------------------------------
 # --- Budgets, batching and spilling ---
 # ---------------------------------------