From 8270e529583dc92ae5ddb67e65a6db49c996a1e7 Mon Sep 17 00:00:00 2001 From: Andrew Tritt Date: Thu, 10 Aug 2023 14:00:37 -0700 Subject: [PATCH] Fix bug when running on GPUs (#43) --- src/gtnet/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gtnet/predict.py b/src/gtnet/predict.py index 92630e1..13c6906 100644 --- 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