Linting

neural_lam/__init__.py

@@ -1,8 +1,10 @@
+# First-party
 import neural_lam.config
 import neural_lam.interaction_net
 import neural_lam.metrics
 import neural_lam.models
 import neural_lam.utils
 import neural_lam.vis
-from .weather_dataset import WeatherDataset
 
+# Local
+from .weather_dataset import WeatherDataset

neural_lam/create_parameter_weights.py

@@ -11,8 +11,7 @@
 from tqdm import tqdm
 
 # Local
-from . import config
-from .weather_dataset import WeatherDataset
+from . import WeatherDataset, config
 
 
 class PaddedWeatherDataset(torch.utils.data.Dataset):

neural_lam/models/__init__.py

@@ -1,5 +1,6 @@
+# Local
+from .base_graph_model import BaseGraphModel
+from .base_hi_graph_model import BaseHiGraphModel
 from .graph_lam import GraphLAM
 from .hi_lam import HiLAM
 from .hi_lam_parallel import HiLAMParallel
-from .base_graph_model import BaseGraphModel
-from .base_hi_graph_model import BaseHiGraphModel

neural_lam/train_model.py

@@ -10,10 +10,8 @@
 from lightning_fabric.utilities import seed
 
 # Local
-from . import config, utils, WeatherDataset
-from .models import GraphLAM
-from .models import HiLAM
-from .models import HiLAMParallel
+from . import WeatherDataset, config, utils
+from .models import GraphLAM, HiLAM, HiLAMParallel
 
 MODELS = {
     "graph_lam": GraphLAM,