implement floyd on tmat_class atsp generation

ai4co · Oct 17, 2024 · b0e4aa0 · b0e4aa0
1 parent f979492
commit b0e4aa0
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/rl4co/envs/routing/atsp/generator.py b/rl4co/envs/routing/atsp/generator.py
@@ -61,11 +61,6 @@ def _generate(self, batch_size) -> TensorDict:
         dms[..., torch.arange(self.num_loc), torch.arange(self.num_loc)] = 0
         log.info("Using TMAT class (triangle inequality): {}".format(self.tmat_class))
         if self.tmat_class:
-            while True:
-                old_dms = dms.clone()
-                dms, _ = (
-                    dms[..., :, None, :] + dms[..., None, :, :].transpose(-2, -1)
-                ).min(dim=-1)
-                if (dms == old_dms).all():
-                    break
+            for i in range(self.num_loc):
+                dms = torch.minimum(dms, dms[..., :, [i]] + dms[..., [i], :])
         return TensorDict({"cost_matrix": dms}, batch_size=batch_size)