# Buggy code
def forward(self, input: Tensor) -> Tensor:  # type: ignore[no-redef] # noqa: F811
    if isinstance(input, Tensor):
        prev_features = [input]
    else:
        prev_features = input

    if self.memory_efficient and self.any_requires_grad(prev_features):
        if torch.jit.is_scripting():
            raise Exception("Memory Efficient not supported in JIT")

        bottleneck_output = self.call_checkpoint_bottleneck(prev_features)
    else:
        bottleneck_output = self.bn_function(prev_features)

    new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))
    if self.drop_rate > 0:
        new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
    return new_features


# 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, input: Tensor) -> Tensor:  # noqa: F811
    if isinstance(input, Tensor):
        prev_features = [input]
    else:
        prev_features = input

    if self.memory_efficient and self.any_requires_grad(prev_features):
        if torch.jit.is_scripting():
            raise Exception("Memory Efficient not supported in JIT")

        bottleneck_output = self.call_checkpoint_bottleneck(prev_features)
    else:
        bottleneck_output = self.bn_function(prev_features)

    new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))
    if self.drop_rate > 0:
        new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
    return new_features