Hello all, suppose I have a scalar function f(t,x) that is a nn.Module:

```
class f(nn.Module):
def __init__(self):
super(f, self).__init__()
self.net = nn.Sequential(
nn.Linear(2, 50),
nn.Tanh(),
nn.Linear(50, 1),
)
def forward(self, t, x):
return self.net(x)
```

I would like to build a **nn.module** that takes the input (t,x) and outputs the gradient of f with respect to x, roughly like below:

```
class dfdx(nn.Module):
def __init__(self):
super(dfdx, self).__init__()
def forward(self, t, x):
return df_dx
```

In this case, x=[x1, x2] is a 2x1 tensor, so df_dx should be a 2x1 tensor [df/dx1, df/dx2].

I found a similar question, and tried the following

```
class f(nn.Module):
def __init__(self):
super(f, self).__init__()
self.net = nn.Sequential(
nn.Linear(2, 50),
nn.Tanh(),
nn.Linear(50, 1),
)
def forward(self, t, x):
return self.net(x)
def compute_u_x(self, t, x):
self.u_x = torch.autograd.functional.jacobian(self, x, create_graph=True)
self.u_x = torch.squeeze(self.u_x)
return self.u_x
```

and hope that

```
func = f()
func.compute_u_x
```

would achieve the goal (Though I am not sure if `func.compute_u_x`

is a **nn.module**). But got the following error:

```
TypeError: forward() missing 1 required positional argument: 'x'
```

I understand why I get this error but have no idea how to fix it. Is there any way to solve the problem? Any help is highly appreciated!