From 256fd1aab0213919b140732b455bf0d1d8a5eb1a Mon Sep 17 00:00:00 2001
From: Brett Kuprel <brkuprel@gmail.com>
Date: Mon, 4 Jul 2022 07:21:54 -0400
Subject: [PATCH] empty cache if cuda available

---
 README.md              | 2 +-
 min_dalle/min_dalle.py | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 7ad8fc6..c588160 100644
--- a/README.md
+++ b/README.md
@@ -81,7 +81,7 @@ display(image)
 Use `image_from_text.py` to generate images from the command line.
 
 ```bash
-$ python image_from_text.py --text='artificial intelligence' --seed=7
+$ python image_from_text.py --text='artificial intelligence' --no-mega --seed=7
 ```
 <img src="https://github.com/kuprel/min-dalle/raw/main/examples/artificial_intelligence.jpg" alt="min-dalle" width="200"/>
 
diff --git a/min_dalle/min_dalle.py b/min_dalle/min_dalle.py
index 9f12858..35509db 100644
--- a/min_dalle/min_dalle.py
+++ b/min_dalle/min_dalle.py
@@ -142,7 +142,6 @@ class MinDalle:
         params = torch.load(self.detoker_params_path)
         self.detokenizer.load_state_dict(params)
         del params
-        torch.cuda.empty_cache()
         if torch.cuda.is_available(): self.detokenizer = self.detokenizer.cuda()
 
 
@@ -176,7 +175,7 @@ class MinDalle:
             encoder_state
         )
         if not self.is_reusable: del self.decoder
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
         return image_tokens