Add lines to force gradient calculation.

During inference only.

casanovo/denovo/model.py

@@ -758,7 +758,8 @@ def validation_step(
             The loss of the validation step.
         """
         # Record the loss.
-        loss = self.training_step(batch, mode="valid")
+        with torch.set_grad_enabled(True):
+            loss = self.training_step(batch, mode="valid")
         if not self.calculate_precision:
             return loss
 
@@ -809,28 +810,29 @@ def predict_step(
             and amino acid-level confidence scores.
         """
         predictions = []
-        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,
+        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,
+                        )
                     )
-                )
 
         return predictions