From 1a4e745d68bbff84b0d9b2cfcbc1f831c433c11c Mon Sep 17 00:00:00 2001
From: ngastzepeda <n.gastzepeda@outlook.com>
Date: Tue, 5 Dec 2023 02:35:27 +0900
Subject: [PATCH] use self.capacity to calculate demand to cover capacities not
 included in CAPACITIES dict

---
 rl4co/envs/routing/cvrp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rl4co/envs/routing/cvrp.py b/rl4co/envs/routing/cvrp.py
index fd6f92a6..45b0ab5a 100644
--- 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,