Fixed logger creation

Author: vfdev-5

README.md

@@ -56,5 +56,6 @@ python -u -m torch.distributed.launch --nproc_per_node=2 main_fixmatch.py model=
 #### 8 TPUs on Colab
 
 ```bash
-python -u main_fixmatch.py model=WRN-28-2 distributed.backend=xla-tpu
-``
+python -u main_fixmatch.py model=resnet18 distributed.backend=xla-tpu distributed.nproc_per_node=8
+# or python -u main_fixmatch.py model=WRN-28-2 distributed.backend=xla-tpu distributed.nproc_per_node=8
+```

main_fixmatch.py

@@ -34,7 +34,12 @@ def unpack_from_tensor(t):
     return sorted_op_names[k_index], bins, error
 
 
-def training(local_rank, cfg, logger):
+def training(local_rank, cfg):
+
+    logger = setup_logger(
+        "FixMatch Training",
+        distributed_rank=idist.get_rank()
+    )
 
     if local_rank == 0:
         logger.info(cfg.pretty())
@@ -176,11 +181,7 @@ def update_cta_rates():
 def main(cfg: DictConfig) -> None:
 
     with idist.Parallel(backend=cfg.distributed.backend, nproc_per_node=cfg.distributed.nproc_per_node) as parallel:
-        logger = setup_logger(
-            "FixMatch Training",
-            distributed_rank=idist.get_rank()
-        )
-        parallel.run(training, cfg, logger)
+        parallel.run(training, cfg)
 
 
 if __name__ == "__main__":

main_fully_supervised.py

@@ -8,7 +8,12 @@
 import trainers
 
 
-def training(local_rank, cfg, logger):
+def training(local_rank, cfg):
+
+    logger = setup_logger(
+        "Fully-Supervised Training",
+        distributed_rank=idist.get_rank()
+    )
 
     if local_rank == 0:
         logger.info(cfg.pretty())
@@ -68,11 +73,7 @@ def train_step(engine, batch):
 def main(cfg: DictConfig) -> None:
 
     with idist.Parallel(backend=cfg.distributed.backend, nproc_per_node=cfg.distributed.nproc_per_node) as parallel:
-        logger = setup_logger(
-            "Fully-Supervised Training",
-            distributed_rank=idist.get_rank()
-        )
-        parallel.run(training, cfg, logger)
+        parallel.run(training, cfg)
 
 
 if __name__ == "__main__":