From 2cb0a25b24d90807b216aa521d25c63d28c6f3bb Mon Sep 17 00:00:00 2001
From: psbang <psb77@cornell.edu>
Date: Mon, 16 Dec 2024 15:35:02 -0500
Subject: [PATCH] Add torchtune to docker image packages (#1456)

---
 kaggle_requirements.txt | 1 +
 tests/test_torchtune.py | 9 +++++++++
 2 files changed, 10 insertions(+)
 create mode 100644 tests/test_torchtune.py

diff --git a/kaggle_requirements.txt b/kaggle_requirements.txt
index 30f819c5..8379d0f2 100644
--- a/kaggle_requirements.txt
+++ b/kaggle_requirements.txt
@@ -140,6 +140,7 @@ tensorflow_decision_forests
 timm
 torchinfo
 torchmetrics
+torchtune
 tsfresh
 vtk
 wandb
diff --git a/tests/test_torchtune.py b/tests/test_torchtune.py
new file mode 100644
index 00000000..b2090ecc
--- /dev/null
+++ b/tests/test_torchtune.py
@@ -0,0 +1,9 @@
+import unittest
+
+import subprocess
+
+class TestTorchtune(unittest.TestCase):
+    def test_help(self):
+        ret_code = subprocess.run(["tune", "--help"])
+        self.assertEqual(0, ret_code.returncode)
+        self.assertIsNone(ret_code.stderr)
\ No newline at end of file