Merge pull request #105 from ai4co/101_capacity_bug

Fix capacity bug in CVRP #101
ai4co · Dec 5, 2023 · ad38abd · ad38abd
2 parents 253cce6 + 1a4e745
commit ad38abd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rl4co/envs/routing/cvrp.py b/rl4co/envs/routing/cvrp.py
@@ -240,7 +240,7 @@ def generate_data(self, batch_size) -> TensorDict:
             {
                 "locs": locs_with_depot[..., 1:, :],
                 "depot": locs_with_depot[..., 0, :],
-                "demand": demand / CAPACITIES[self.num_loc],
+                "demand": demand / self.capacity,
                 "capacity": capacity,
             },
             batch_size=batch_size,