Add utility to inspect a model's parameters (to get dtype/device)

1 files changed, 19 insertions, 0 deletions

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