Skip to content

Commit 8270e52

Browse files
authored
Fix bug when running on GPUs (#43)
1 parent 3f1278d commit 8270e52

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

src/gtnet/predict.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ def run_torchscript_inference(fastas, model, conf_models, window, step, vocab, s
171171
s = e
172172

173173
del total_chunks
174-
total_chunks = torch.tensor(tmp_chunks)
174+
total_chunks = torch.tensor(tmp_chunks, device=device)
175175
del tmp_chunks
176176

177177
del all_levels_aggregated

0 commit comments

Comments
 (0)