always use dict.__getitem__ in dict/OrderedDict iteration (#676)

thunder/core/interpreter.py

@@ -2345,7 +2345,7 @@ def __iter__(self):
         return self
 
     def __next__(self):
-        return self._mapping[next(self._key_iter)]
+        return dict.__getitem__(self._mapping, next(self._key_iter))
 
 
 class MappingValuesWrapper:
@@ -2369,7 +2369,7 @@ def __iter__(self):
 
     def __next__(self):
         k = next(self._key_iter)
-        return k, self._mapping[k]
+        return k, dict.__getitem__(self._mapping, k)
 
 
 class MappingItemsWrapper:

thunder/tests/test_interpreter.py

@@ -3214,3 +3214,19 @@ def test_litgpt(jit):
     result = jfn(*args, **kwargs)
 
     assert_close(result, fn(*args, **kwargs))
+
+
+def test_transformer_model_output():
+    pytest.importorskip("transformers")
+    from transformers.utils.generic import ModelOutput
+
+    def fn(x):
+        mo = ModelOutput(foo=x)
+        return mo["foo"]
+
+    x = torch.randn(3)
+    expected = fn(x)
+
+    actual = thunder.jit(fn)(x)
+
+    assert expected is actual