2025-02-20 nightly release (62c0740)

pytorch · Feb 20, 2025 · 1f788ab · 1f788ab
1 parent f419ab8
commit 1f788ab
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrec/distributed/planner/proposers.py b/torchrec/distributed/planner/proposers.py
@@ -913,7 +913,7 @@ def allocate_budget(
             )
             actual_increase_bytes = new_cache_size_bytes - cache_size_bytes
 
-            budget -= torch.sum(actual_increase_bytes)
+            budget -= torch.sum(actual_increase_bytes).item()  # pyre-ignore[58]
             cache_size_bytes = new_cache_size_bytes
             # TODO: consider trade off of using remaining budget to push >0.95 tables
             # to HBM vs spending that budget on improving hit rate on other tables in
@@ -930,7 +930,7 @@ def allocate_budget(
                 logger.debug(
                     f"[allocate_budget] {promotes.sum()} tables exceeded ceiling, promoting to save {budget_reclaimed} bytes"
                 )
-                budget += budget_reclaimed
+                budget += budget_reclaimed  # pyre-ignore[58]
                 # force these tables to 1.0 to ensure promotion
                 cache_size_bytes[promotes] = max_cache_size_bytes[promotes]