From 9646dfc0e6b24453d5fbe2baedbac709cc56699e Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 26 Jul 2022 09:47:54 -0700
Subject: [PATCH] fix path_or_state bug

---
 dalle2_pytorch/trainer.py | 2 +-
 dalle2_pytorch/version.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/dalle2_pytorch/trainer.py b/dalle2_pytorch/trainer.py
index 41ce286..87ed3b5 100644
--- a/dalle2_pytorch/trainer.py
+++ b/dalle2_pytorch/trainer.py
@@ -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)
 
diff --git a/dalle2_pytorch/version.py b/dalle2_pytorch/version.py
index 858de17..da2182f 100644
--- a/dalle2_pytorch/version.py
+++ b/dalle2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.0.5'
+__version__ = '1.0.6'