def train_nn(model, data, num_epoch=5000):
    train_dataset = TensorDataset(torch.Tensor(data.Xtrain), torch.Tensor(data.Ytrain))
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
    test_dataset = TensorDataset(torch.Tensor(data.Xtest), torch.Tensor(data.Ytest))
    test_dataloader = DataLoader(dataset=test_dataset, batch_size=128)
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    losses = []
    for epoch in range(num_epoch):
        for (inputs, targets) in train_dataloader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
        model.eval()
        loss = 0.0
        for (inputs, targets) in test_dataloader:
            outputs = model(inputs)
            loss += criterion(outputs, targets).data
        losses.append((loss.data // len(test_dataloader)))
        model.train()
    return losses
def train_nn(model, data, num_epoch=5000):
    train_dataset = TensorDataset(torch.Tensor(data.Xtrain), torch.Tensor(data.Ytrain))
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
    test_dataset = TensorDataset(torch.Tensor(data.Xtest), torch.Tensor(data.Ytest))
    test_dataloader = DataLoader(dataset=test_dataset, batch_size=128)
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    losses = []
    for epoch in range(num_epoch):
        for (inputs, targets) in train_dataloader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
        model.eval()
        loss = 0.0
        for (inputs, targets) in test_dataloader:
            outputs = model(inputs)
            loss += criterion(outputs, targets).data
        losses.append((loss.data // len(test_dataloader)))
        model.train()
    return losses