Solving an ode using pytorch

I am trying to solve an ode using pytorch. The ode has the form

du/dt = cos(2*3.14*t)

I parameterise my neural network as a two layer linear network
with tanh as an activation function in between. The layer takes in 1 dimensional input and returns 1 dimensional output with hidden layer size being 32.

def f(x):
    function that computes the output of the neural net
    l1 = torch.matmul(W1.T, x).reshape(-1, 1)
    l1_act = leakyrelu(l1)
    l2 = torch.matmul(W2.T, l1_act)
    return l2

def g(t):
    the form of the solution chosen such that boundary condition
    u(0) = 1. is satisfied 
    return t*f(t) + torch.tensor([1.])

def loss(t, eps):
    the loss function which is simply the loss 
    on the gradient of the net being equal to the analytical gradient 
    return torch.mean(((g(t+eps) - g(t)) /(eps) - torch.cos(torch.tensor([2*3.14*t])))**2)

These are the model definition

x = torch.tensor([.01])
eps = torch.tensor([0.000345])
T0 = 0
T1 = 1
nsamples = 100
t = torch.linspace(T0, T1, nsamples)

W1 = Variable(torch.ones(1, 32), requires_grad=True)
W2 = Variable(torch.ones(32, 1), requires_grad = True)

I generate a 100 datapoints between 0 and 1 and train the network on these datapoints for 5000 epochs. My training loop looks something like this

learning_rate = 1e-3
for it in tqdm(range(5000)):
    err = 0
    for ti in t:
        ti = torch.tensor([ti])
        err += loss(ti, eps)
    err = err / nsamples
    err.backward() -= learning_rate * -= learning_rate *
    if it%100==0:
        grad_w1 =
        grad_w2 =
        fig, a= plt.subplots(1, 2, figsize =(6, 3))

What I notice is that the gradient updates go to 0 and the network doesnt manage to learn the values. I had been following along in an example in Julia and there the code seems to work. I was wondering if there is something wrong in my specification. If someone can point me in the right direction it would be great. This is my second attempt to solve. I tried using the nn.Module along with nn.Linear layer along with Adam optimizer and there too the error seems to be the same