diff --git a/lambeq/training/nelder_mead_optimizer.py b/lambeq/training/nelder_mead_optimizer.py
index 6b2377c0..c2f78eb6 100644
--- a/lambeq/training/nelder_mead_optimizer.py
+++ b/lambeq/training/nelder_mead_optimizer.py
@@ -259,7 +259,7 @@ def objective(self, x: Iterable[Any], y: ArrayLike, w: ArrayLike) -> float:
                 raise ValueError(
                     'Objective function must return a scalar'
                 ) from e
-        return result   # type: ignore[return-value]
+        return result
 
     def backward(self, batch: tuple[Iterable[Any], np.ndarray]) -> float:
         """Calculate the gradients of the loss function.
diff --git a/lambeq/training/pytorch_trainer.py b/lambeq/training/pytorch_trainer.py
index 3d8bf100..f70e4e04 100644
--- a/lambeq/training/pytorch_trainer.py
+++ b/lambeq/training/pytorch_trainer.py
@@ -174,7 +174,7 @@ def validation_step(
         with torch.no_grad():
             y_hat = self.model(x)
             loss = self.loss_function(y_hat, y.to(self.device))
-        return y_hat, loss.item()
+        return y_hat.detach(), loss.detach().item()
 
     def training_step(
             self,
@@ -196,8 +196,9 @@ def training_step(
         x, y = batch
         y_hat = self.model(x)
         loss = self.loss_function(y_hat, y.to(self.device))
-        self.train_costs.append(loss.item())
         self.optimizer.zero_grad()
         loss.backward()
         self.optimizer.step()
-        return y_hat, loss.item()
+        loss_item = loss.detach().item()
+        self.train_costs.append(loss_item)
+        return y_hat.detach(), loss_item