[BugFix] Fix compile during _check_keys

ghstack-source-id: cfed094
Pull Request resolved: #1239

(cherry picked from commit 2ad9f95)

Author: Vincent Moens

tensordict/utils.py

@@ -1768,7 +1768,8 @@ def _check_keys(
         is_leaf=_is_leaf_nontensor,
     )
     # TODO: compile doesn't like set() over an arbitrary object
-    if is_compiling():
+    is_comp = is_compiling()
+    if is_comp:
         keys_set = {k for k in keys}  # noqa: C416
     else:
         keys_set: set[str] = set(keys)
@@ -1781,7 +1782,7 @@ def _check_keys(
         if not strict:
             keys_set = keys_set.intersection(k)
         else:
-            if is_compiling():
+            if is_comp:
                 k = {v for v in k}  # noqa: C416
             else:
                 k = set(k)
@@ -1790,7 +1791,10 @@ def _check_keys(
                     f"got keys {keys} and {set(td.keys())} which are incompatible"
                 )
     if strict:
-        return list(keys)
+        if is_comp:
+            return [key for key in keys]  # noqa: C416
+        else:
+            return list(keys)
     return keys_set