I am trying to train a simple neural network for regression problem. Here X_train, y_train, X_valid, y_vlid are my X/y numpy arrays for training/validation.  After running the training, my training and validation loss are empty (NaN).  Here is my
model, criterion and training function, is there a problem with my function?
class Net_Reg(nn.Module):
    def __init__(self):
        super(Net_Reg, self).__init__()
        self.fc1 = nn.Linear(15, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


reg_model = Net_Reg()
criterion = torch.nn.MSELoss()
optimizer = optim.SGD(reg_model.parameters(), lr=0.01)


X_trn = torch.from_numpy(X_train).float()
y_trn = torch.from_numpy(y_train).float()
t_ds = TensorDataset(X_trn, y_trn)

X_val = torch.from_numpy(X_valid).float()
y_val = torch.from_numpy(y_valid).float()
v_ds = TensorDataset(X_trn, y_trn)

def reg_train(model, t_ds, v_ds, bs, epochs):
    
    t_loader = DataLoader(t_ds, batch_size=bs, shuffle=True)
    v_loader = DataLoader(v_ds, batch_size=bs, shuffle=True)
    
    epoch_count = []
    train_loss = []
    val_loss = []
    
    for epoch in range(epochs):
        
        model.train()
        
        train_running_loss = 0
        
        for t_inputs, t_labels in t_loader:
            
            t_outputs = model(t_inputs)
            t_loss = criterion(t_outputs, t_labels)
            
            train_running_loss += t_loss.item()
  
            optimizer.zero_grad()
            t_loss.backward()
            optimizer.step()
        
        epoch_count.append(epoch)
        train_loss.append(train_running_loss)
        
        model.eval()
        
        val_running_loss = 0
        
        for v_inputs, v_labels in v_loader:
            
            v_outputs = model(v_inputs) 
            v_loss = criterion(v_outputs, v_labels)
            
            val_running_loss += v_loss.item()
        
        val_loss.append(val_running_loss)

    print ("training loss: {:.4f}\nvalidation loss: {:.4f}".format(sum(train_loss)/len(t_ds),
                                                                sum(val_loss)/len(v_ds)))
    plt.plot(epoch_count, train_loss, c='b')
    plt.plot(epoch_count, val_loss, c='r')
    
return model