#Credits - https://github.com/hunkim/PyTorchZeroToAll/blob/master/05_linear_regression.py

import torch
from torch.autograd import Variable

x_data = Variable(torch.Tensor([[1.0],[2.0],[3.0]]))
y_data = Variable(torch.Tensor(([2.0],[4.0],[6.0])))

class Model(torch.nn.Module):
    def __init__(self):
        #Constructor
        super(Model,self).__init__()
        self.linear = torch.nn.Linear(1,1) #One in and One out
        
    def forward(self,x):
        #Variable of input data
        #Variable of output data
        y_pred = self.linear(x)
        return y_pred
    
#Initialize model
model = Model()

#Loss Function and Optimizer
criterion = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(),lr=0.01)
total_loss = 0
#Training Loop
for epoch in range(500):
    y_pred = model(x_data)
    
    #compute and print loss
    loss = criterion(y_pred,y_data)
    print(epoch,loss.item())
    
    #Zero gradients
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
#After training
hour_var = Variable(torch.Tensor([[4.0]]))
y_pred = model(hour_var)


print('Predict after training',4,model(hour_var).data[0][0])