From b516158a238537b03534ee1453f325fa140e47ec Mon Sep 17 00:00:00 2001
From: Laurin Luttmann <laurin.luttmann@gmail.com>
Date: Thu, 12 Sep 2024 10:08:59 +0200
Subject: [PATCH] [Fix] fixed JSSP import issues

---
 rl4co/envs/__init__.py            | 9 ++++-----
 rl4co/envs/scheduling/__init__.py | 9 ++++-----
 rl4co/envs/scheduling/jssp/env.py | 2 +-
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/rl4co/envs/__init__.py b/rl4co/envs/__init__.py
index f3b65b7d..1c43a7d1 100644
--- a/rl4co/envs/__init__.py
+++ b/rl4co/envs/__init__.py
@@ -4,6 +4,9 @@
 # EDA
 from rl4co.envs.eda import DPPEnv, MDPPEnv
 
+# Graph
+from rl4co.envs.graph import FLPEnv, MCPEnv
+
 # Routing
 from rl4co.envs.routing import (
     ATSPEnv,
@@ -25,11 +28,7 @@
 )
 
 # Scheduling
-from rl4co.envs.scheduling import FFSPEnv, FJSPEnv, SMTWTPEnv
-from rl4co.envs.scheduling.jssp.env import JSSPEnv
-
-# Graph
-from rl4co.envs.graph import MCPEnv, FLPEnv
+from rl4co.envs.scheduling import FFSPEnv, FJSPEnv, JSSPEnv, SMTWTPEnv
 
 # Register environments
 ENV_REGISTRY = {
diff --git a/rl4co/envs/scheduling/__init__.py b/rl4co/envs/scheduling/__init__.py
index 40b5571e..febb183c 100644
--- a/rl4co/envs/scheduling/__init__.py
+++ b/rl4co/envs/scheduling/__init__.py
@@ -1,5 +1,4 @@
-from rl4co.envs.scheduling.ffsp.env import FFSPEnv
-from rl4co.envs.scheduling.fjsp.env import FJSPEnv
-
-# from rl4co.envs.scheduling.jssp.env import JSSPEnv
-from rl4co.envs.scheduling.smtwtp.env import SMTWTPEnv
+from .ffsp.env import FFSPEnv
+from .fjsp.env import FJSPEnv
+from .jssp.env import JSSPEnv
+from .smtwtp.env import SMTWTPEnv
diff --git a/rl4co/envs/scheduling/jssp/env.py b/rl4co/envs/scheduling/jssp/env.py
index 702ceda7..2381836a 100644
--- a/rl4co/envs/scheduling/jssp/env.py
+++ b/rl4co/envs/scheduling/jssp/env.py
@@ -4,7 +4,7 @@
 from tensordict import TensorDict
 from torch._tensor import Tensor
 
-from rl4co.envs import FJSPEnv
+from rl4co.envs.scheduling.fjsp.env import FJSPEnv
 from rl4co.utils.ops import gather_by_index
 
 from .generator import JSSPFileGenerator, JSSPGenerator