make test happy

ai4co · Jan 12, 2025 · 99a05f3 · 99a05f3
1 parent dcd5821
commit 99a05f3
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/rl4co/envs/routing/cvrp/env.py b/rl4co/envs/routing/cvrp/env.py
@@ -12,7 +12,10 @@
 
 from .generator import CVRPGenerator
 
-from .local_search import local_search
+try:
+    from .local_search import local_search
+except:  # In case when we fail to build HGS
+    local_search = None
 from .render import render
 
 log = get_pylogger(__name__)
@@ -253,7 +256,7 @@ def replace_selected_actions(
     def local_search(td: TensorDict, actions: torch.Tensor, **kwargs) -> torch.Tensor:
         assert (
             local_search is not None
-        ), "Cannot import local_search module. Check if `pyvrp` is installed."
+        ), "Cannot import local_search module. Check if HGS is built."
         return local_search(td, actions, **kwargs)
 
     @staticmethod