Fix for device issue (#126)

* Will's fix for device issue * more descriptive GPU choices
matsengrp · Mar 6, 2025 · 92d29c9 · 92d29c9
1 parent 25a3a56
commit 92d29c9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/netam/common.py b/netam/common.py
@@ -207,15 +207,15 @@ def find_least_used_cuda_gpu(mem_round_val=1300):
 
     # Check prioritization order:
     if max(utilization) > 0:
-        print("used utilization")
+        print("GPU chosen via utilization")
         return utilization.index(min(utilization))  # Least utilized GPU
 
     if max(memory_used) > 0:
-        print("used memory")
+        print("GPU chosen via memory")
         return memory_used.index(min(memory_used))  # Least memory used GPU
 
     if len(set(gpu_counts.values())) > 1:
-        print("used processes")
+        print("GPU chosen via process count")
         return min(
             uuid_to_index[uuid]
             for uuid, count in gpu_counts.items()

diff --git a/netam/models.py b/netam/models.py
@@ -649,7 +649,7 @@ def selection_factors_of_aa_str(self, aa_sequence: Tuple[str, str]) -> Tensor:
             ).squeeze(0)
         else:
             result = result.squeeze(0)
-        return split_heavy_light_model_outputs(result, idx_seq.squeeze(0))
+        return split_heavy_light_model_outputs(result.cpu(), idx_seq.squeeze(0).cpu())
 
 
 class TransformerBinarySelectionModelLinAct(AbstractBinarySelectionModel):