jsonify modules better

Author: JasonGross

gbmi/utils/hashing.py

@@ -27,6 +27,7 @@
 
 import numpy
 import torch
+import transformer_lens
 from transformer_lens import HookedTransformer
 
 # Implemented for https://github.com/lemon24/reader/issues/179
@@ -149,6 +150,16 @@ def _json_default(
             exclude_filter=exclude_filter,
             dictify_by_default=dictify_by_default,
         )
+    elif isinstance(thing, torch.nn.Module):
+        return _json_dumps(
+            {
+                "type": type(thing),
+                "repr": repr(thing),
+                "module.parameters": list(thing.parameters()),
+            },
+            exclude_filter=exclude_filter,
+            dictify_by_default=dictify_by_default,
+        )
     elif isinstance(thing, type):
         return f"{thing.__module__}.{thing.__name__}"
     elif (