Add unit test

tests/test_decorator.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+
+"""Trivial version test."""
+
+import unittest
+
+import torch
+
+from torch_max_mem import maximize_memory_utilization
+
+
+def knn(x, y, batch_size, k: int = 3):
+    return torch.cat(
+        [
+            torch.cdist(x[start: start + batch_size], y).topk(k=k, dim=0, largest=False).indices
+            for start in range(0, x.shape[0], batch_size)
+        ],
+        dim=0,
+    )
+
+
+wrapped_knn = maximize_memory_utilization(parameter_name="batch_size")(knn)
+
+
+class TestDecorator(unittest.TestCase):
+    """Test the decorator."""
+
+    def test_knn(self):
+        x = torch.rand(100, 100)
+        y = torch.rand(200, 100)
+        for batch_size in [1, 10, x.shape[0]]:
+            self.assertEqual(knn(x, y, batch_size), wrapped_knn(x, y, batch_size=x.shape[0]))