From 1b3716049007723a0d2bcdb0fd91ba29ff430a1e Mon Sep 17 00:00:00 2001
From: Philipp Emanuel Weidmann <pew@worldwidemann.com>
Date: Sun, 21 Sep 2025 16:04:41 +0530
Subject: [PATCH] Fix model loading issues

---
 config.default.toml  | 2 +-
 src/heretic/model.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/config.default.toml b/config.default.toml
index 16fb39d..5becae6 100644
--- a/config.default.toml
+++ b/config.default.toml
@@ -1,4 +1,4 @@
-dtypes = ["float32", "float16"]
+dtypes = ["float32", "float16", "bfloat16"]
 
 device_map = "auto"
 
diff --git a/src/heretic/model.py b/src/heretic/model.py
index 5a3169b..ffb582e 100644
--- a/src/heretic/model.py
+++ b/src/heretic/model.py
@@ -47,6 +47,7 @@ class Model:
                 self.generate([settings.test_prompt], max_new_tokens=1)
             except Exception as error:
                 self.model = None
+                empty_cache()
                 print(f"[red]Failed[/] ({error})")
                 continue