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

1 files changed, 17 insertions, 0 deletions

diff --git a/modules/torch_utils.py b/modules/torch_utils.py
new file mode 100644
index 00000000..e5b52393
--- /dev/null
+++ b/modules/torch_utils.py
@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import torch.nn
+
+
+def get_param(model) -> torch.nn.Parameter:
+    """
+    Find the first parameter in a model or module.
+    """
+    if hasattr(model, "model") and hasattr(model.model, "parameters"):
+        # Unpeel a model descriptor to get at the actual Torch module.
+        model = model.model
+
+    for param in model.parameters():
+        return param
+
+    raise ValueError(f"No parameters found in model {model!r}")