add process group for distillation teacher checkpoint loader (#248)

cijose · web-flow · commit 0e35cfdea1f8 · 2025-10-21T14:55:57.000+02:00
diff --git a/dinov3/train/ssl_meta_arch.py b/dinov3/train/ssl_meta_arch.py
@@ -339,6 +339,7 @@ def init_weights(self) -> None:
                     self.cfg.distillation.checkpoint_path,
                     skip_load_keys=["dino_loss.center", "ibot_patch_loss.center"],
                     keys_not_sharded=["backbone.rope_embed.periods", "qkv.bias_mask"],
+                    process_group=distributed.get_default_process_group(),
                 )
             else:
                 logger.info("Init teacher to distil from, used for testing purpose only")

Original file line number	Diff line number	Diff line change
`@@ -339,6 +339,7 @@ def init_weights(self) -> None:`
`339`	`339`	`self.cfg.distillation.checkpoint_path,`
`340`	`340`	`skip_load_keys=["dino_loss.center", "ibot_patch_loss.center"],`
`341`	`341`	`keys_not_sharded=["backbone.rope_embed.periods", "qkv.bias_mask"],`
	`342`	`+ process_group=distributed.get_default_process_group(),`
`342`	`343`	`)`
`343`	`344`	`else:`
`344`	`345`	`logger.info("Init teacher to distil from, used for testing purpose only")`