# Buggy code
def forward(self, init_features: Tensor) -> Tensor:  # type: ignore[override]
    features = [init_features]
    for name, layer in self.items():
        new_features = layer(features)
        features.append(new_features)
    return torch.cat(features, 1)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "repro.py", line 7, in <module>
torch.jit.script(model).save('densenet161.pt')
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 911, in script
return torch.jit._recursive.create_script_module(
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 370, in create_script_module
return create_script_module_impl(nn_module, concrete_type, stubs_fn)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 388, in _construct
init_fn(script_module)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 406, in init_fn
scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 388, in _construct
init_fn(script_module)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 406, in init_fn
scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 382, in create_script_module_impl
method_stubs = stubs_fn(nn_module)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 618, in infer_methods_to_compile
stubs.append(make_stub_from_method(nn_module, method))
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 52, in make_stub_from_method
return make_stub(func, method_name)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 37, in make_stub
ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/frontend.py", line 259, in get_jit_def
return build_def(ctx, fn_def, type_line, def_name, self_name=self_name)
File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/frontend.py", line 288, in build_def
type_comment_decl = torch._C.parse_type_comment(type_line)
RuntimeError: expected type comment but found 'def' here:
def forward(self, init_features: Tensor) -> Tensor:  # type: ignore[override]
~~~ <--- HERE

# Fixed code
def forward(self, init_features: Tensor) -> Tensor:
    features = [init_features]
    for name, layer in self.items():
        new_features = layer(features)
        features.append(new_features)
    return torch.cat(features, 1)