Pytorch Sinusoid estimation

import torch
from torch.autograd import Variable
import torch.nn.functional as F
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
from functools import reduce

%matplotlib inline

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.hn_1 = 4
        self.hn_2 = 16

        self.gru_1 = torch.nn.GRUCell(1, self.hn_1, bias=True)
        self.gru_2 = torch.nn.GRUCell(self.hn_1, self.hn_2, bias=True)
        self.linear_1 = torch.nn.Linear(self.hn_2, 1)
        self.h_1 = None
        self.h_2 = None

    def reset_hidden_state(self):
        self.h_1 = Variable(torch.randn(1, self.hn_1))
        self.h_2 = Variable(torch.randn(1, self.hn_2))

    def forward(self, x):
        gru_outputs = []
        for i in range(x.shape[0] - 1):
            self.h_1 = self.gru_1(x[i].view(1, 1), self.h_1)
            self.h_2 = self.gru_2(self.h_1, self.h_2)
            linear_out = self.linear_1(self.h_2)
            gru_outputs.append(linear_out)

        return gru_outputs


def generate_sin(mesh):
    func = np.sin(mesh)*5
    noise = np.random.standard_normal(size=func.shape[0])
    return func# + noise

dataset_source = [
    np.arange(np.random.randint(1, 5),  np.random.randint(6, 20), 0.1).astype(np.float32).reshape(-1, 1) for i in range(100)
]

dataset_target = [
    generate_sin(i.squeeze()).astype(np.float32).reshape(-1, 1) for i in dataset_source
]

num_epochs = 100

neu = Model()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(neu.parameters(), lr=0.0001)


for epoch in range(num_epochs):
    is_last_call = False
    for count, x in enumerate(zip(dataset_source, dataset_target)):
        x_source, x_target = x
        neu.reset_hidden_state()
        optimizer.zero_grad()

        x_source = Variable(torch.from_numpy(x_source))
        x_target = Variable(torch.from_numpy(x_target))

        outputs = neu(x_source)
        loss = 0
        for i in range(len(outputs)):
            loss += criterion(outputs[i], x_target[i])
        loss.backward(retain_graph=False)
        optimizer.step()

        if count == len(dataset) - 1:
            is_first_call = True
    print('loss on epoch {}: {}'.format(epoch, loss.data.numpy()))

test_x = np.arange(0, 10, 0.1).astype(np.float32).reshape(-1, 1)
test_y = generate_sin(test_x.squeeze()).astype(np.float32).reshape(-1, 1)
test_pred = np.array([i.data.numpy().squeeze() for i in neu(Variable(torch.from_numpy(test_x)))])

plt.figure(figsize=(15, 5))
plt.plot(test_y.squeeze())
plt.plot(test_pred.squeeze())
plt.show()