Fixed formatting

axonn-ai · Nov 1, 2023 · 77e4036 · 77e4036
1 parent 704168c
commit 77e4036
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/axonn/communication.py b/axonn/communication.py
@@ -31,17 +31,17 @@ def __init__(
         """
         self.world_rank = MPI.COMM_WORLD.Get_rank()
         self.world_size = MPI.COMM_WORLD.Get_size()
-        
+
         if gpus_per_node is None:
-            self.backend = 'gloo'
+            self.backend = "gloo"
             self.is_gpu_available = False
             self.local_rank = os.environ.get("SLURM_LOCALID", 0)
         else:
-            self.backend = 'nccl'
+            self.backend = "nccl"
             self.is_gpu_available = True
             self.local_rank = self.world_rank % gpus_per_node
             torch.cuda.set_device(self.local_rank)
-        
+
         G_intra = G_intra_r * G_intra_c
         assert (
             G_inter * G_data * G_intra == self.world_size