fix: isort fails

ecmwf · May 30, 2024 · 6584f0e · 6584f0e
1 parent 64e35dd
commit 6584f0e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py
@@ -10,12 +10,12 @@
 import uuid
 
 import torch
+from anemoi.utils.config import DotConfig
 from hydra.utils import instantiate
 from torch_geometric.data import HeteroData
 
 from anemoi.models.models.encoder_processor_decoder import AnemoiModelEncProcDec
 from anemoi.models.preprocessing import Processors
-from anemoi.utils.config import DotConfig
 
 
 class AnemoiModelInterface(torch.nn.Module):

diff --git a/src/anemoi/models/layers/block.py b/src/anemoi/models/layers/block.py
@@ -491,7 +491,7 @@ def forward(
 
         query, key, value, edges = self.shard_qkve_heads(query, key, value, edges, shapes, batch_size, model_comm_group)
 
-        # TODO: Is this alright?
+        # TODO: remove magic number
         num_chunks = self.num_chunks if self.training else 4  # reduce memory for inference
 
         if num_chunks > 1:

diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py
@@ -12,6 +12,7 @@
 
 import einops
 import torch
+from anemoi.utils.config import DotConfig
 from hydra.utils import instantiate
 from torch import Tensor
 from torch import nn
@@ -21,7 +22,6 @@
 
 from anemoi.models.distributed.shapes import get_shape_shards
 from anemoi.models.layers.graph import TrainableTensor
-from anemoi.utils.config import DotConfig
 
 LOGGER = logging.getLogger(__name__)