tweaks

matsengrp · Jan 24, 2025 · d12c1cd · d12c1cd
1 parent 9111b8b
commit d12c1cd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/netam/common.py b/netam/common.py
@@ -542,6 +542,7 @@ def parallelize_function(
     max_worker_count = min(mp.cpu_count() // 2, max_workers)
     if max_worker_count <= 1:
         return function
+    force_spawn()
 
     @wraps(function)
     def wrapper(*args, **kwargs):

diff --git a/netam/framework.py b/netam/framework.py
@@ -349,6 +349,7 @@ def crepe_exists(prefix):
 
 def trimmed_shm_model_outputs_of_crepe(crepe, parents):
     """Model outputs trimmed to the length of the parent sequences."""
+    crepe.to("cpu")
     rates, csp_logits = parallelize_function(crepe)(parents)
     rates = rates.cpu().detach()
     csps = torch.softmax(csp_logits, dim=-1).cpu().detach()