[fix] nn c++ : segfault in modulelist and moduledict (pytorch#93074)

Fixes pytorch#73565

Pull Request resolved: pytorch#93074
Approved by: https://github.com/albanD

Author: kshitij12345

test/cpp/api/moduledict.cpp

@@ -299,3 +299,11 @@ TEST_F(ModuleDictTest, PrettyPrintModuleDict) {
       "  (lstm): torch::nn::LSTM(input_size=4, hidden_size=5, num_layers=1, bias=true, batch_first=false, dropout=0, bidirectional=false)\n"
       ")");
 }
+
+TEST_F(ModuleDictTest, InvalidAt) {
+  torch::OrderedDict<std::string, std::shared_ptr<Module>> ordereddict = {
+      {"linear", Linear(10, 3).ptr()}};
+  ModuleDict dict(ordereddict);
+  ASSERT_THROWS_WITH(
+      dict->at<torch::nn::Dropout2dImpl>("linear"), "Unable to cast module");
+}

test/cpp/api/modulelist.cpp

@@ -300,3 +300,9 @@ TEST_F(ModuleListTest, RangeBasedForLoop) {
     module->pretty_print(buffer);
   }
 }
+
+TEST_F(ModuleListTest, InvalidAt) {
+  torch::nn::ModuleList m(torch::nn::Linear(1, 2));
+  ASSERT_THROWS_WITH(
+      m->at<torch::nn::Dropout2dImpl>(0), "Unable to cast module");
+}

torch/csrc/api/include/torch/nn/modules/container/moduledict.h

@@ -178,7 +178,14 @@ class ModuleDictImpl : public Cloneable<ModuleDictImpl> {
     static_assert(
         torch::detail::is_module<T>::value,
         "Can only call ModuleList::at with an nn::Module type");
-    return *modules_[key]->as<T>();
+    auto module = modules_[key]->as<T>();
+    TORCH_CHECK(
+        module,
+        "Unable to cast module[",
+        key,
+        "] to ",
+        c10::demangle(typeid(T).name()));
+    return *module;
   }
 
   /// Attempts to return the module at the given key as the requested type.
@@ -189,7 +196,14 @@ class ModuleDictImpl : public Cloneable<ModuleDictImpl> {
     static_assert(
         torch::detail::is_module<T>::value,
         "Can only call ModuleList::at with an nn::Module type");
-    return *modules_[key]->as<T>();
+    const auto module = modules_[key]->as<T>();
+    TORCH_CHECK(
+        module,
+        "Unable to cast module[",
+        key,
+        "] to ",
+        c10::demangle(typeid(T).name()));
+    return *module;
   }
 
   /// Removes and returns the `Module` associated with the given `key`.

torch/csrc/api/include/torch/nn/modules/container/modulelist.h

@@ -147,7 +147,14 @@ class ModuleListImpl : public Cloneable<ModuleListImpl> {
         torch::detail::is_module<T>::value,
         "Can only call ModuleList::at with an nn::Module type");
     TORCH_CHECK(index < size(), "Index out of range");
-    return *modules_[index]->as<T>();
+    auto module = modules_[index]->as<T>();
+    TORCH_CHECK(
+        module,
+        "Unable to cast module[",
+        index,
+        "] to ",
+        c10::demangle(typeid(T).name()));
+    return *module;
   }
 
   /// Attempts to return the module at the given index as the requested type.
@@ -159,7 +166,14 @@ class ModuleListImpl : public Cloneable<ModuleListImpl> {
         torch::detail::is_module<T>::value,
         "Can only call ModuleList::at with an nn::Module type");
     TORCH_CHECK(index < size(), "Index out of range");
-    return *modules_[index]->as<T>();
+    const auto module = modules_[index]->as<T>();
+    TORCH_CHECK(
+        module,
+        "Unable to cast module[",
+        index,
+        "] to ",
+        c10::demangle(typeid(T).name()));
+    return *module;
   }
 
   /// Attempts to return a `std::shared_ptr` whose dynamic type is that of the