Support NNModules as dict keys (#116723)

Summary:
X-link: pytorch/pytorch#116723
Approved by: https://github.com/lezcano

Reviewed By: osalpekar

Differential Revision: D52648982

Pulled By: voznesenskym

fbshipit-source-id: 55c94bbda1e16e3ed20ce61e15c184e0a90bcea4

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -1059,8 +1059,11 @@ def iter_contains(items, search, tx, check_tensor_identity=False):
     return found
 
 
-def tensor_to_id(value):
-    return [id(k) if isinstance(k, torch.Tensor) else k for k in value.keys()]
+def tensor_or_module_to_id(value):
+    return [
+        id(k) if isinstance(k, (torch.Tensor, torch.nn.Module)) else k
+        for k in value.keys()
+    ]
 
 
 def const_repr(x, *, local) -> str: