From 5768afc776a66bb94e77a9c1daebeea58fa731d5 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Sun, 31 Dec 2023 00:20:30 +0200
Subject: Add utility to inspect a model's parameters (to get dtype/device)

---
 test/test_torch_utils.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 test/test_torch_utils.py

(limited to 'test')

diff --git a/test/test_torch_utils.py b/test/test_torch_utils.py
new file mode 100644
index 00000000..f1aec832
--- /dev/null
+++ b/test/test_torch_utils.py
@@ -0,0 +1,19 @@
+import types
+
+import pytest
+import torch
+
+from modules.torch_utils import get_param
+
+
+@pytest.mark.parametrize("wrapped", [True, False])
+def test_get_param(wrapped):
+    mod = torch.nn.Linear(1, 1)
+    cpu = torch.device("cpu")
+    mod.to(dtype=torch.float16, device=cpu)
+    if wrapped:
+        # more or less how spandrel wraps a thing
+        mod = types.SimpleNamespace(model=mod)
+    p = get_param(mod)
+    assert p.dtype == torch.float16
+    assert p.device == cpu
-- 
cgit v1.2.3