From 7b2b8be0e7b5619e89afc85f5db2204a804ff04d Mon Sep 17 00:00:00 2001
From: Dave Berenbaum <dave@iterative.ai>
Date: Wed, 10 Jul 2024 13:25:46 -0400
Subject: [PATCH] make hf tests work

---
 pyproject.toml                       | 2 +-
 tests/frameworks/test_huggingface.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index a35a76b..a245029 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,7 +57,7 @@ tests = [
   "pytest_voluptuous",
   "dpath",
   "transformers[torch]",
-  "keras<3"
+  "tf-keras"
 ]
 dev = [
   "dvclive[all,tests]",
diff --git a/tests/frameworks/test_huggingface.py b/tests/frameworks/test_huggingface.py
index 36a4187..18b3990 100644
--- a/tests/frameworks/test_huggingface.py
+++ b/tests/frameworks/test_huggingface.py
@@ -103,6 +103,7 @@ def args():
         num_train_epochs=2,
         save_strategy="epoch",
         report_to="none",  # Disable auto-reporting to avoid duplication
+        use_cpu=True,
     )