[BugFix] Macbook related bug

ai4co · Aug 22, 2023 · 305e1fa · 305e1fa
1 parent 7d47178
commit 305e1fa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/rl4co/models/zoo/active_search/search.py b/rl4co/models/zoo/active_search/search.py
@@ -98,7 +98,7 @@ def setup(self, stage="fit"):
         )
         self.instance_rewards = torch.zeros(dataset_size)
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int) -> int | None:
+    def on_train_batch_start(self, batch: Any, batch_idx: int):
         """Called before training (i.e. search) for a new batch begins.
         We re-load the original policy state dict and configure the optimizer.
         """

diff --git a/rl4co/models/zoo/eas/search.py b/rl4co/models/zoo/eas/search.py
@@ -106,7 +106,7 @@ def setup(self, stage="fit"):
         self.instance_solutions = []
         self.instance_rewards = []
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int) -> int | None:
+    def on_train_batch_start(self, batch: Any, batch_idx: int):
         """Called before training (i.e. search) for a new batch begins.
         We re-load the original policy state dict and configure all parameters not to require gradients.
         We do the rest in the training step.