diff --git a/CHANGELOG.md b/CHANGELOG.md
index 628da137..08b77443 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Changed
 
 - The CLI has been overhauled to use subcommands.
-- Upgraded to Lightning >=2.0
+- Upgraded to Lightning >=2.0.
 - Checkpointing is configured to save the top-k models instead of all.
 - Log steps rather than epochs as units of progress during training.
 - Validation performance metrics are logged (and added to tensorboard) at the validation epoch, and training loss is logged at the end of training epoch, i.e. training and validation metrics are logged asynchronously.
@@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Nicely format logged warnings.
 - `every_n_train_steps` has been renamed to `val_check_interval` in accordance to the corresponding Pytorch Lightning parameter.
 - Training batches are randomly shuffled.
+- Upgraded to Torch >=2.1.
 
 ### Removed
 
@@ -36,7 +37,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Correctly refer to input peak files by their full file path.
 - Specifying custom residues to retrain Casanovo is now possible.
 - Upgrade to depthcharge v0.2.3 to fix sinusoidal encoding and for the `PeptideTransformerDecoder` hotfix.
-- Enable gradients during prediction and validation to avoid NaNs from occuring as a temporary workaround until a new Pytorch version is available.
 - Correctly report amino acid precision and recall during validation.
 
 ## [3.3.0] - 2023-04-04
diff --git a/casanovo/denovo/model.py b/casanovo/denovo/model.py
index 5851aa76..26766ead 100644
--- a/casanovo/denovo/model.py
+++ b/casanovo/denovo/model.py
@@ -751,9 +751,7 @@ def validation_step(
             The loss of the validation step.
         """
         # Record the loss.
-        # FIXME: Temporary workaround to avoid the NaN bug.
-        with torch.set_grad_enabled(True):
-            loss = self.training_step(batch, mode="valid")
+        loss = self.training_step(batch, mode="valid")
         if not self.calculate_precision:
             return loss
 
@@ -804,30 +802,28 @@ def predict_step(
             and amino acid-level confidence scores.
         """
         predictions = []
-        # FIXME: Temporary workaround to avoid the NaN bug.
-        with torch.set_grad_enabled(True):
-            for (
-                precursor_charge,
-                precursor_mz,
-                spectrum_i,
-                spectrum_preds,
-            ) in zip(
-                batch[1][:, 1].cpu().detach().numpy(),
-                batch[1][:, 2].cpu().detach().numpy(),
-                batch[2],
-                self.forward(batch[0], batch[1]),
-            ):
-                for peptide_score, aa_scores, peptide in spectrum_preds:
-                    predictions.append(
-                        (
-                            spectrum_i,
-                            precursor_charge,
-                            precursor_mz,
-                            peptide,
-                            peptide_score,
-                            aa_scores,
-                        )
+        for (
+            precursor_charge,
+            precursor_mz,
+            spectrum_i,
+            spectrum_preds,
+        ) in zip(
+            batch[1][:, 1].cpu().detach().numpy(),
+            batch[1][:, 2].cpu().detach().numpy(),
+            batch[2],
+            self.forward(batch[0], batch[1]),
+        ):
+            for peptide_score, aa_scores, peptide in spectrum_preds:
+                predictions.append(
+                    (
+                        spectrum_i,
+                        precursor_charge,
+                        precursor_mz,
+                        peptide,
+                        peptide_score,
+                        aa_scores,
                     )
+                )
 
         return predictions
 
diff --git a/pyproject.toml b/pyproject.toml
index 2b836a41..551954ab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
     "scikit-learn",
     "spectrum_utils",
     "tensorboard",
-    "torch>=2.0",
+    "torch>=2.1",
     "tqdm",
 ]
 dynamic = ["version"]