From d0dbdb3f6ec05905d2748ee9661cc47b53a50a6a Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 19 Jan 2021 15:07:40 -0800
Subject: [PATCH] be able to specify a seed

---
 big_sleep/big_sleep.py | 10 +++++++++-
 big_sleep/cli.py       |  2 ++
 big_sleep/version.py   |  2 +-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py
index ffc4b06..8ff0fca 100644
--- a/big_sleep/big_sleep.py
+++ b/big_sleep/big_sleep.py
@@ -33,6 +33,9 @@ def signal_handling(signum,frame):
 
 # helpers
 
+def exists(val):
+    return val is not None
+
 def open_folder(path):
     if os.path.isfile(path):
         path = os.path.dirname(path)
@@ -165,9 +168,14 @@ def __init__(
         iterations = 1050,
         save_progress = False,
         bilinear = False,
-        open_folder = True
+        open_folder = True,
+        seed = None
     ):
         super().__init__()
+
+        if exists(seed):
+            torch.manual_seed(seed)
+
         self.epochs = epochs
         self.iterations = iterations
 
diff --git a/big_sleep/cli.py b/big_sleep/cli.py
index 8253836..aaf513b 100644
--- a/big_sleep/cli.py
+++ b/big_sleep/cli.py
@@ -13,6 +13,7 @@ def train(
     overwrite = False,
     save_progress = False,
     bilinear = False,
+    seed = None,
     open_folder = True
 ):
 
@@ -26,6 +27,7 @@ def train(
         save_every = save_every,
         save_progress = save_progress,
         bilinear = bilinear,
+        seed = seed,
         open_folder = open_folder
     )
 
diff --git a/big_sleep/version.py b/big_sleep/version.py
index 44b1806..407b8a2 100644
--- a/big_sleep/version.py
+++ b/big_sleep/version.py
@@ -1 +1 @@
-__version__ = '0.2.6'
+__version__ = '0.2.7'