MAINT: attempt to limit memory leaks (#89)

* limit queries searched * del pytorch variables
stsievert · Nov 25, 2020 · e452970 · e452970
1 parent 90be880
commit e452970
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/salmon/triplets/algs/_adaptive_runners.py b/salmon/triplets/algs/_adaptive_runners.py
@@ -134,7 +134,7 @@ def get_queries(
         rng = None
         if random_state:
             rng = check_random_state(random_state)
-        for pwr in itertools.count(start=13):
+        for pwr in range(12, 20 + 1):
             queries, scores = self.search.score(num=2 ** pwr, random_state=rng)
             ret_queries.append(queries)
             ret_scores.append(scores)

diff --git a/salmon/triplets/algs/adaptive/_embed.py b/salmon/triplets/algs/adaptive/_embed.py
@@ -198,6 +198,8 @@ def partial_fit(self, answers, sample_weight=None):
             self.meta_["model_updates"] += 1
             logger.info("%s", self.meta_)
             if self.meta_["num_grad_comps"] >= eg_deadline:
+                del loss
+                del losses
                 break
         return self