mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Fix distributed flag bug w/ flex device handling

Browse Source
This commit is contained in:
Ross Wightman 2024-02-03 16:26:15 -08:00
parent bee0471f91
commit a08b57e801
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
timm/utils/distributed.py
View File

@ -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