[Performance] Minor efficiency improvements (#703)

pytorch · Mar 24, 2024 · 82e83af · 82e83af
1 parent 8ea3f13
commit 82e83af
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/tensordict/_td.py b/tensordict/_td.py
@@ -520,6 +520,14 @@ def __eq__(self, other: object) -> T | bool:
             keys1 = set(self.keys())
             keys2 = set(other.keys())
             if len(keys1.difference(keys2)) or len(keys1) != len(keys2):
+                keys1 = sorted(
+                    keys1,
+                    key=lambda key: "".join(key) if isinstance(key, tuple) else key,
+                )
+                keys2 = sorted(
+                    keys2,
+                    key=lambda key: "".join(key) if isinstance(key, tuple) else key,
+                )
                 raise KeyError(f"keys in tensordicts mismatch, got {keys1} and {keys2}")
             d = {}
             for key, item1 in self.items():

diff --git a/tensordict/base.py b/tensordict/base.py
@@ -5252,8 +5252,9 @@ def _register_tensor_class(cls):
 
 
 def _is_tensor_collection(datatype):
-    out = _TENSOR_COLLECTION_MEMO.get(datatype, None)
-    if out is None:
+    try:
+        out = _TENSOR_COLLECTION_MEMO[datatype]
+    except KeyError:
         if issubclass(datatype, TensorDictBase):
             out = True
         elif _is_tensorclass(datatype):