empty cache if cuda available

2022-07-04 07:21:54 -04:00
parent 9602055eaa
commit 256fd1aab0
2 changed files with 2 additions and 3 deletions
--- 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