YannAhlgrim/wilds-ijepa
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

DDP guard if single GPUgio

Browse Source
This commit is contained in:
YannAhlgrim
2026-05-16 11:52:10 +02:00
parent bad00f306e
commit 79fc9564ea
1 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/train.py
+5
View File
@@ -214,11 +214,16 @@ def main(args, resume_preempt=False):
ipe_scale=ipe_scale, ipe_scale=ipe_scale,
use_bfloat16=use_bfloat16, use_bfloat16=use_bfloat16,
) )
if dist.is_available() and dist.is_initialized() and world_size > 1:
encoder = DistributedDataParallel(encoder, static_graph=True) encoder = DistributedDataParallel(encoder, static_graph=True)
predictor = DistributedDataParallel(predictor, static_graph=True) predictor = DistributedDataParallel(predictor, static_graph=True)
target_encoder = DistributedDataParallel(target_encoder) target_encoder = DistributedDataParallel(target_encoder)
for p in target_encoder.parameters(): for p in target_encoder.parameters():
p.requires_grad = False p.requires_grad = False
else:
logger.info("Distributed process group not initialized; running without DDP")
for p in target_encoder.parameters():
p.requires_grad = False
# -- momentum schedule # -- momentum schedule
momentum_scheduler = ( momentum_scheduler = (