Lint fix

Author: pomonam

algorithmic_efficiency/checkpoint_utils.py

@@ -119,7 +119,9 @@ def maybe_restore_checkpoint(framework: str,
 
   else:
     checkpoint_state = latest_ckpt
-    if isinstance(model_params, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)):
+    if isinstance(
+        model_params,
+        (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)):
       model_params = model_params.module
     model_params.load_state_dict(checkpoint_state['model_params'])
     checkpoint_state['model_params'] = model_params
@@ -196,7 +198,9 @@ def save_checkpoint(framework: str,
     opt_state = jax.device_get(jax_utils.unreplicate(opt_state))
     model_state = jax.device_get(jax_utils.unreplicate(model_state))
   else:
-    if isinstance(model_params, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)):
+    if isinstance(
+        model_params,
+        (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)):
       model_params = model_params.module
     model_params = model_params.state_dict()
     optimizer_state_dict = {}

submission_runner.py

@@ -319,8 +319,8 @@ def train_once(
     metrics_logger = None
     if RANK == 0:
       metrics_logger = logger_utils.set_up_loggers(log_dir,
-                                                     flags.FLAGS,
-                                                     hyperparameters)
+                                                   flags.FLAGS,
+                                                   hyperparameters)
       workload.attach_metrics_logger(metrics_logger)
 
   global_start_time = get_time()
@@ -470,13 +470,13 @@ def train_once(
   metrics = {'eval_results': eval_results, 'global_step': global_step}
 
   if log_dir is not None and RANK == 0:
-      metrics_logger.append_scalar_metrics(
-          {'score': train_state['accumulated_submission_time']},
-          global_step=global_step,
-          preemption_count=preemption_count)
-      metrics_logger.finish()
-      if save_checkpoints:
-        checkpoint_utils.save_checkpoint(
+    metrics_logger.append_scalar_metrics(
+        {'score': train_state['accumulated_submission_time']},
+        global_step=global_step,
+        preemption_count=preemption_count)
+    metrics_logger.finish()
+    if save_checkpoints:
+      checkpoint_utils.save_checkpoint(
           framework=FLAGS.framework,
           optimizer_state=optimizer_state,
           model_params=model_params,