fix path_or_state bug

This commit is contained in:
Phil Wang
2022-07-26 09:47:54 -07:00
parent 62043acb2f
commit 9646dfc0e6
2 changed files with 2 additions and 2 deletions

View File

@@ -300,7 +300,7 @@ class DiffusionPriorTrainer(nn.Module):
# all processes need to load checkpoint. no restriction here
if isinstance(path_or_state, str):
path = Path(path)
path = Path(path_or_state)
assert path.exists()
loaded_obj = torch.load(str(path), map_location=self.device)