Hi, I tried, if I do:
print("loss.grad_fn: ",loss.grad_fn)
I get this:
loss.grad_fn:  <MseLossBackward object at 0x7f9152ca7358>
My code is the same as the one above in which I changed return x with return output
However, the code is below for convenience
I arrive to see only
Debug 5
loss.grad_fn:  <MseLossBackward object at 0x7f9152905b00>

# In[15]:


import torch
import numpy as np

from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split

import matplotlib
import matplotlib.pyplot as plt
import numpy as np


# In[16]:


LR = 1e-6
MAX_EPOCH = 10
BATCH_SIZE = 512

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")


# In[17]:


class SineApproximator(nn.Module):
    def __init__(self):
        super(SineApproximator, self).__init__()
        self.regressor = nn.Sequential(nn.Linear(1, 1024),
                                       nn.ReLU(inplace=True),
                                       nn.Linear(1024, 1024),
                                       nn.ReLU(inplace=True),
                                       nn.Linear(1024, 1))
    def forward(self, x):
        output = self.regressor(x)
        return output

X = np.random.rand(10**5) * 2 * np.pi
y = np.sin(X)

print(X)
print(np.size(X))
print(y)
print(np.size(y))


# In[18]:



# # Data for plotting
# # t = np.arange(0.0, 2.0, 0.01)
# # s = 1 + np.sin(2 * np.pi * t)

# matplotlib.rcParams['agg.path.chunksize'] = 10000

# fig, ax = plt.subplots()
# ax.plot(X, y)

# # ax.set(xlabel='X', ylabel='y',
# #        title='Test plot')
# # ax.grid()

# # fig.savefig("test.png")
# plt.show()


# In[19]:


class SineApproximator(nn.Module):
    def __init__(self):
        super(SineApproximator, self).__init__()
        self.regressor = nn.Sequential(nn.Linear(1, 1024),
                                       nn.ReLU(inplace=True),
                                       nn.Linear(1024, 1024),
                                       nn.ReLU(inplace=True),
                                       nn.Linear(1024, 1))
    def forward(self, x):
        output = self.regressor(x)
        return output

X = np.random.rand(10**5) * 2 * np.pi
y = np.sin(X)

# Data for plotting
# t = np.arange(0.0, 2.0, 0.01)
# s = 1 + np.sin(2 * np.pi * t)

fig, ax = plt.subplots()
ax.plot(X, y)

ax.set(xlabel='X', ylabel='y',
       title='Test plot')
ax.grid()

# fig.savefig("test.png")
plt.show()

X_train, X_val, y_train, y_val = map(torch.tensor, train_test_split(X, y, test_size=0.2))
train_dataloader = DataLoader(TensorDataset(X_train.unsqueeze(1), y_train.unsqueeze(1)), batch_size=BATCH_SIZE,
                              pin_memory=True, shuffle=True)
val_dataloader = DataLoader(TensorDataset(X_val.unsqueeze(1), y_val.unsqueeze(1)), batch_size=BATCH_SIZE,
                            pin_memory=True, shuffle=True)

model = SineApproximator().to(device)
optimizer = optim.Adam(model.parameters(), lr=LR)
criterion = nn.MSELoss(reduction="mean")

# training loop
train_loss_list = list()
val_loss_list = list()
print("Debug 1")
for epoch in range(MAX_EPOCH):
    print("epoch %d / %d" % (epoch+1, MAX_EPOCH))
    model.train()
    # training loop
    temp_loss_list = list()
    for X_train, y_train in train_dataloader:
        X_train = X_train.type(torch.float32).to(device)
        y_train = y_train.type(torch.float32).to(device)
        print("Debug 2")
        optimizer.zero_grad()
        print("Debug 3")
        score = model(X_train)
        print("Debug 4")
        loss = criterion(input=score, target=y_train)
        print("Debug 5")
        print("loss.grad_fn: ",loss.grad_fn)
        loss.requires_grad = True
        print("Debug 6")
        loss.backward()
        
        optimizer.step()

        temp_loss_list.append(loss.detach().cpu().numpy())
    
    temp_loss_list = list()
    for X_train, y_train in train_dataloader:
        X_train = X_train.type(torch.float32).to(device)
        y_train = y_train.type(torch.float32).to(device)

        score = model(X_train)
        loss = criterion(input=score, target=y_train)

        temp_loss_list.append(loss.detach().cpu().numpy())
    
    train_loss_list.append(np.average(temp_loss_list))

    # validation
    model.eval()
    
    temp_loss_list = list()
    for X_val, y_val in val_dataloader:
        X_val = X_val.type(torch.float32).to(device)
        y_val = y_val.type(torch.float32).to(device)

        score = model(X_val)
        loss = criterion(input=score, target=y_val)

        temp_loss_list.append(loss.detach().cpu().numpy())
    
    val_loss_list.append(np.average(temp_loss_list))

    print("\ttrain loss: %.5f" % train_loss_list[-1])
    print("\tval loss: %.5f" % val_loss_list[-1])

Thanks for your help!