Fix

SIAnalytics · Dec 1, 2022 · 65de9b9 · 65de9b9
1 parent f542291
commit 65de9b9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 5 deletions.
diff --git a/siatune/mm/tasks/base.py b/siatune/mm/tasks/base.py
@@ -122,7 +122,6 @@ def context_aware_run(self, searched_cfg: Dict) -> Any:
         context = dict(
             args=deepcopy(self.args),
             searched_cfg=deepcopy(ImmutableContainer.decouple(searched_cfg)),
-            # checkpoint_dir=checkpoint_dir,
         )
         return context_manager(self.run)(**context)
 

diff --git a/siatune/mm/tasks/mmseg.py b/siatune/mm/tasks/mmseg.py
@@ -80,7 +80,6 @@ def build_model(self,
                 The train opt. Defaults to None.
             test_cfg (Optional[Config]):
                 The Test opt. Defaults to None.
-
         Returns:
             torch.nn.Module: The model.
         """
@@ -98,7 +97,6 @@ def build_dataset(
             cfg (Config): The configs.
             default_args (Optional[Config]):
                 The default args. Defaults to None.
-
         Returns:
             torch.utils.data.Dataset: The dataset.
         """

diff --git a/siatune/mm/tasks/mmtrainbase.py b/siatune/mm/tasks/mmtrainbase.py
@@ -4,7 +4,7 @@
 import mmcv
 import torch
 from ray.air.config import ScalingConfig
-from ray.train.torch import TorchTrainer
+from ray.train.torch import TorchConfig, TorchTrainer
 
 from .base import BaseTask
 from .builder import TASKS
@@ -71,4 +71,5 @@ def create_trainable(self) -> TorchTrainer:
                 use_gpu=True,
                 resources_per_worker=dict(
                     CPU=self.num_cpus_per_worker,
-                    GPU=self.num_gpus_per_worker)))
+                    GPU=self.num_gpus_per_worker)),
+            torch_config=TorchConfig(backend='gloo'))