def run(f):
    print(f"{f.__name__}")
    print("-" * len(f.__name__))
    f()
    print()

@run
def test_heaviside():
    class HeavisideModule(nn.Module):
        def __init__(self):
            super().__init__()
        def forward(self,x):
            input,values=x
            return torch.heaviside(input,values)
    
    exported=fx.export_and_import(HeavisideModule(),(torch.randn(1,2),torch.randn(1)),output_type='torch')
    print(exported)