[JIT] module compilation fails when multiple instances of a class have a different type #32469
Description
🐛 Bug
The following example fails:
import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn
model = maskrcnn_resnet50_fpn()
# The following modifications cause the last command to fail
model.roi_heads.mask_roi_pool = None
model.roi_heads.mask_head = None
model.roi_heads.mask_predictor = None
model_script = torch.jit.script(model)
Throws the following:
~/miniconda3/lib/python3.7/site-packages/torch/jit/__init__.py in script(obj, optimize, _frames_up, _rcb)
1253
1254 if isinstance(obj, torch.nn.Module):
-> 1255 return torch.jit._recursive.recursive_script(obj)
1256
1257 qualified_name = _qualified_name(obj)
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in recursive_script(nn_module)
532 check_module_initialized(nn_module)
533
--> 534 return create_script_module(nn_module, infer_methods_to_compile(nn_module))
535
536 def try_compile_fn(fn, loc):
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs)
291 """
292 check_module_initialized(nn_module)
--> 293 concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)
294 cpp_module = torch._C._create_module_with_type(concrete_type.jit_type)
295
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in get_or_create_concrete_type(self, nn_module)
234 return nn_module._concrete_type
235
--> 236 concrete_type_builder = infer_concrete_type_builder(nn_module)
237
238 nn_module_type = type(nn_module)
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in infer_concrete_type_builder(nn_module)
119 else:
120 # otherwise we get the concrete module type for item and add it to concrete_type
--> 121 sub_concrete_type = concrete_type_store.get_or_create_concrete_type(item)
122 concrete_type_builder.add_module(name, sub_concrete_type)
123
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in get_or_create_concrete_type(self, nn_module)
234 return nn_module._concrete_type
235
--> 236 concrete_type_builder = infer_concrete_type_builder(nn_module)
237
238 nn_module_type = type(nn_module)
~/miniconda3/lib/python3.7/site-packages/torch/jit/_recursive.py in infer_concrete_type_builder(nn_module)
116 if attr_type is not None:
117 # if the type can be inferred, it should be a module interface type
--> 118 sub_concrete_type = torch._C.ConcreteModuleType.from_jit_type(attr_type)
119 else:
120 # otherwise we get the concrete module type for item and add it to concrete_type
RuntimeError: type->cast<ClassType>() INTERNAL ASSERT FAILED at /pytorch/torch/csrc/jit/script/concrete_module_type.cpp:44, please report a bug to PyTorch. (fromJitType at /pytorch/torch/csrc/jit/script/concrete_module_type.cpp:44)
I think this is something to do with module type sharing, i tried to debug a bit but couldn't get anywhere. cc @suo