Skip to content

Commit

Permalink
Move N-term tokens to the correct device
Browse files Browse the repository at this point in the history
The device changes between initialization and actual running, hence doing this in __init__ is insufficient.
  • Loading branch information
bittremieux committed Dec 7, 2023
1 parent 38212f5 commit bc56c49
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions casanovo/denovo/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down

0 comments on commit bc56c49

Please sign in to comment.