distribute

Wonters · Wonters · commit 5a023dc2c6e7 · 2025-04-22T20:05:35.000+02:00
diff --git a/src/ml.py b/src/ml.py
@@ -474,7 +474,7 @@ def load_checkpoint(self):
                     map_location={"cuda:0": "mps", "cuda": "mps"},
                 )
             else:
-                checkpoint = torch.load(self.checkpoint + "/model.pth")
+                checkpoint = torch.load(self.checkpoint + "/model.pth",map_location=f"cuda:{local_rank}")
             embedding_weights = {
                 k: v
                 for k, v in checkpoint.items()

Original file line number	Diff line number	Diff line change
`@@ -474,7 +474,7 @@ def load_checkpoint(self):`
`474`	`474`	`map_location={"cuda:0": "mps", "cuda": "mps"},`
`475`	`475`	`)`
`476`	`476`	`else:`
`477`		`- checkpoint = torch.load(self.checkpoint + "/model.pth")`
	`477`	`+ checkpoint = torch.load(self.checkpoint + "/model.pth",map_location=f"cuda:{local_rank}")`
`478`	`478`	`embedding_weights = {`
`479`	`479`	`k: v`
`480`	`480`	`for k, v in checkpoint.items()`