Fix bug when running on GPUs (#43)

exabiome · Aug 10, 2023 · 8270e52 · 8270e52
1 parent 3f1278d
commit 8270e52
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/gtnet/predict.py b/src/gtnet/predict.py
@@ -171,7 +171,7 @@ def run_torchscript_inference(fastas, model, conf_models, window, step, vocab, s
             s = e
 
         del total_chunks
-        total_chunks = torch.tensor(tmp_chunks)
+        total_chunks = torch.tensor(tmp_chunks, device=device)
         del tmp_chunks
 
         del all_levels_aggregated