print->logger.debug

Future-House · Nov 25, 2024 · 5e571d9 · 5e571d9
1 parent 4c706b6
commit 5e571d9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ldp/alg/optimizer/replay_buffers.py b/ldp/alg/optimizer/replay_buffers.py
@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import random
 from collections import UserList
 from collections.abc import Awaitable, Callable, Iterator
@@ -10,6 +11,8 @@
 from ldp.graph import eval_mode
 from ldp.graph.async_torch import AsyncTorchModule
 
+logger = logging.getLogger(__name__)
+
 
 class CircularReplayBuffer(UserList[dict]):
     def resize(self, size: int):
@@ -160,6 +163,6 @@ def batched_iter(
 
             # DEBUG
             selected_priorities = prio[sampled_idcs]
-            print(f"Average priority: {selected_priorities.mean()}")
+            logger.debug(f"Average priority: {selected_priorities.mean()}")
 
         return self._batched_iter(buffer, batch_size, shuffle, infinite)