From bc56c49b322824ab448a422d72d4c7e88e9946a9 Mon Sep 17 00:00:00 2001
From: Wout Bittremieux <wout@bittremieux.be>
Date: Thu, 7 Dec 2023 10:54:00 +0100
Subject: [PATCH] Move N-term tokens to the correct device

The device changes between initialization and actual running, hence doing this in __init__ is insufficient.
---
 casanovo/denovo/model.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/casanovo/denovo/model.py b/casanovo/denovo/model.py
index 9e5a34f7..909bac00 100644
--- a/casanovo/denovo/model.py
+++ b/casanovo/denovo/model.py
@@ -268,6 +268,7 @@ def beam_search_decode(
         scores = einops.rearrange(scores, "B L V S -> (B S) L V")
 
         # The main decoding loop.
+        self.n_term = self.n_term.to(self.decoder.device)
         for step in range(0, self.max_length):
             # Terminate beams exceeding the precursor m/z tolerance and track
             # all finished beams (either terminated or stop token predicted).