Merge pull request #1225 from bghira/bugfix/flux-multinode-rank-detec…

…tion

flux: use rank 0 for h100 detection since that is the most realistic setup

helpers/models/flux/transformer.py

@@ -210,11 +210,8 @@ def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
 
         processor = FluxAttnProcessor2_0()
         if torch.cuda.is_available():
-            rank = (
-                torch.distributed.get_rank()
-                if torch.distributed.is_initialized()
-                else 0
-            )
+            # let's assume that the box only ever has H100s.
+            rank = 0
             primary_device = torch.cuda.get_device_properties(rank)
             if primary_device.major == 9 and primary_device.minor == 0:
                 if is_flash_attn_available: