empty cache if cuda available

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"/>
 

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