fix reset

lucidrains · Feb 10, 2021 · cbce951 · cbce951
1 parent 708ac98
commit cbce951
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py
@@ -291,6 +291,7 @@ def set_text(self, text):
 
     def reset(self):
         self.model.reset()
+        self.model = self.model.cuda()
         self.optimizer = Adam(self.model.model.latents.parameters(), self.lr)
 
     def train_step(self, epoch, i, pbar=None):

diff --git a/big_sleep/version.py b/big_sleep/version.py
@@ -1 +1 @@
-__version__ = '0.4.10'
+__version__ = '0.4.11'