Merge remote-tracking branch 'origin/main' into polynet

rl4co/envs/routing/op/env.py

@@ -139,7 +139,7 @@ def _reset(
                     device=device,
                 ),
                 "current_total_prize": torch.zeros(
-                    *batch_size, 1, dtype=torch.float, device=device
+                    *batch_size, dtype=torch.float, device=device
                 ),
                 "i": torch.zeros(
                     (*batch_size,), dtype=torch.int64, device=device