Problem solved:
Input and Targer were not being indexed correctly during the training:
for i in range(num_iterations):
    target = y_tensor.to(device).to(torch.float)
    x_data = x_tensor.to(device).to(torch.float)
    for idx, x in x_data:
        pred = model(x).squeeze(-1)
        optimiser.zero_grad()
        loss = loss_fn(pred,target[idx])
        loss.backward()
        optimiser.step()
    if (i+1) % 1 == 0:
        print("[iteration %03d] loss: %.4f" % (i+1, loss.item()))