def python_only_fn(self, input):
                 # This function won't be compiled, so any
                 # Python APIs can be used
                 import pdb
                 pdb.set_trace()
         def forward(self, input):
             if self.training:
                 self.python_only_fn(input)
             return input * 99

     scripted_module = torch.jit.script(MyModule())
     print(scripted_module.some_entry_point(torch.randn(2, 2)))
     print(scripted_module(torch.randn(2, 2)))

  Example ( Annotating forward of nn.Module using example_inputs):

     import torch
     import torch.nn as nn
     from typing import NamedTuple

     class MyModule(NamedTuple):
     result: List[int]

     class TestNNModule(torch.nn.Module):
         def forward(self, a) -> MyModule:
             result = MyModule(result=a)
             return result

     pdt_model = TestNNModule()
