From a08b57e801ab936ee7a7f54ac605f6257a5ee32f Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Sat, 3 Feb 2024 16:26:15 -0800
Subject: [PATCH] Fix distributed flag bug w/ flex device handling

---
 timm/utils/distributed.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/timm/utils/distributed.py b/timm/utils/distributed.py
index 95655f2c..92b8a6b8 100644
--- a/timm/utils/distributed.py
+++ b/timm/utils/distributed.py
@@ -92,6 +92,7 @@ def init_distributed_device(args):
     args.world_size = result['world_size']
     args.rank = result['global_rank']
     args.local_rank = result['local_rank']
+    args.distributed = args.world_size > 1
     device = torch.device(args.device)
     return device