def METHOD_NAME(self):
    # same as above, but also add a CallFunction in between.
    @torch.jit.ignore
    def python_fn(x):
        return x.relu()
    def fn(a, b, c):
        r = a.relu().relu()
        z = python_fn(r)
        return torch.relu(r), torch.nn.functional.gelu(r), c.cos().relu()
    fn_s = torch.jit.script(fn)
    a = torch.rand((10, 10), requires_grad=False)
    b = torch.rand((10, 10), requires_grad=False)
    c = torch.rand((10, 10), requires_grad=True)
    for i in range(4):
        x_s, y_s, z_s = fn_s(a, b, c)
        x, y, z = fn(a, b, c)
        self.assertEqual(x_s.requires_grad, x.requires_grad)
        self.assertEqual(y_s.requires_grad, y.requires_grad)
        self.assertEqual(z_s.requires_grad, z.requires_grad)