Fix `toIValue` dict iteration (#26856)
Summary:
It's not clear why, but using the `[]` accessory in pybind was causing an error (for something I ran into while scripting MaskR-CNN)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/26856
Pulled By: driazati
Differential Revision: D17589635
fbshipit-source-id: 142a9b4e30bcf2b65f1ba03c2d6775f2baff43f5
diff --git a/torch/csrc/jit/pybind_utils.h b/torch/csrc/jit/pybind_utils.h
index 083b22e..e649322 100644
--- a/torch/csrc/jit/pybind_utils.h
+++ b/torch/csrc/jit/pybind_utils.h
@@ -321,14 +321,14 @@
}
inline IValue createGenericDict(
- py::handle obj,
+ py::dict obj,
const TypePtr& key_type,
const TypePtr& value_type) {
c10::impl::GenericDict elems(key_type, value_type);
elems.reserve(py::len(obj));
- for (auto key : obj) {
+ for (auto entry : obj) {
elems.insert(
- toIValue(key, key_type), toIValue(obj[key], value_type));
+ toIValue(entry.first, key_type), toIValue(entry.second, value_type));
}
return IValue(std::move(elems));
}
@@ -445,7 +445,7 @@
case TypeKind::DictType: {
const auto& dict_type = type->expect<DictType>();
return createGenericDict(
- obj, dict_type->getKeyType(), dict_type->getValueType());
+ py::cast<py::dict>(obj), dict_type->getKeyType(), dict_type->getValueType());
}
case TypeKind::OptionalType: {
// check if it's a none obj since optional accepts NoneType