small edits

casanovo/denovo/model.py

@@ -969,6 +969,7 @@ def predict_step(self, batch, *args):
                 )
         return batch_res
 
+    @torch.compile
     def smart_batch_gen(self, batch):
         all_psm = []
         enc = self.encoder(batch[0])
@@ -1014,14 +1015,13 @@ def on_predict_epoch_end(self, results) -> None:
         results = np.array(results, dtype=object).squeeze()
         with open(self.out_writer.filename, "a") as out_f:
             csv_writer = csv.writer(out_f)
-            for batch in results:
-                for index, t_or_d, peptide, score, per_aa_scores in list(
-                    zip(*batch)
-                ):
-                    csv_writer.writerow(
-                        (index, peptide, bool(t_or_d), score, per_aa_scores)
-                    )
-            out_f.close()
+            for index, t_or_d, peptide, score, per_aa_scores in list(
+                zip(*results)
+            ):
+                csv_writer.writerow(
+                    (index, peptide, bool(t_or_d), score, per_aa_scores)
+                )
+        out_f.close()
 
 
 def calc_match_score(

casanovo/denovo/model_runner.py

@@ -394,7 +394,7 @@ def db_search(
     abs_experiment_dirpath = "/net/noble/vol2/home/vananth3/2023_vananth_denovo-dbsearch/results/2023-08-21_speedup"
     profiler = SimpleProfiler(
         dirpath=abs_experiment_dirpath,
-        filename="casanovo_plasmodium_vanilla",
+        filename="casanovo_plasmodium_batch",
         extended=True,
     )
     trainer = pl.Trainer(