MLP for regression using PyTorch

Question

I am trying to implement a MLP with 3 hidden layers for a regression task. My code runs but the loss is always "nan" so something is wrong and have not been able to figure it out. Here is my implementaiton:

#Import 
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import torch.optim as optim

#Prepare data
X_train_tensor = torch.tensor(train_data_top.drop(columns=['DATE','permno','RET']).values, dtype=torch.float32)
y_train_tensor = torch.tensor(train_data_top['RET'].values, dtype=torch.float32)

X_test_tensor = torch.tensor(test_data_top.drop(columns=['DATE','permno','RET']).values, dtype=torch.float32)
y_test_tensor = torch.tensor(test_data_top['RET'].values, dtype=torch.float32)

X_val_tensor = torch.tensor(val_data_top.drop(columns=['DATE','permno','RET']).values, dtype=torch.float32)
y_val_tensor = torch.tensor(val_data_top['RET'].values, dtype=torch.float32)

train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
val_dataset = TensorDataset(X_val_tensor, y_val_tensor)


batch_size = 256
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

#MLP Model
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size1, hidden_size2, hidden_size3, output_size):
        super(MLP, self).__init__()
        #Create fully connected layers
        self.fc1 = nn.Linear(input_size, hidden_size1)
        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
        self.fc3 = nn.Linear(hidden_size2, hidden_size3)
        self.fc4 = nn.Linear(hidden_size3, output_size)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

# Adjust input_size based on the number of features in your input data
input_size = train_dataset[0][0].shape[0]
print(input_size)
hidden_size1, hidden_size2, hidden_size3 = 256, 128, 64
output_size = 1  # For regression
learning_rate = 0.0001

model = MLP(input_size, hidden_size1, hidden_size2, hidden_size3, output_size)
criterion = nn.MSELoss()  # Mean Squared Error loss for regression
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

num_epochs = 10
#Training
for epoch in range(num_epochs):
    model.train()
    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()


    #Validation
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for val_batch_x, val_batch_y in val_loader:
            val_outputs = model(val_batch_x)
            val_loss += criterion(val_outputs, val_batch_y).item()

    val_loss /= len(val_loader)

    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}, Val Loss: {val_loss:.4f}')

print('Training finished.')

# R-squared function
def calculate_r_squared(y_true, y_pred):
    return r2_score(y_true, y_pred)

# Testing
model.eval()
test_loss = 0.0
all_predictions = []
all_targets = []

with torch.no_grad():
    for test_batch_x, test_batch_y in test_loader:
        test_outputs = model(test_batch_x)
        test_loss += criterion(test_outputs, test_batch_y).item()

        all_predictions.extend(test_outputs.numpy())
        all_targets.extend(test_batch_y.numpy())

test_loss /= len(test_loader)

# Calculate R-squared
r_squared = calculate_r_squared(all_targets, all_predictions)

print(f'Test Loss: {test_loss:.4f}')
print(f'R-squared: {r_squared:.4f}')

print('Training and Testing finished.')

All my data is normalized and any rosws with null values have been dropped. The output for my epoch for loop always prints the loss at "nan"

I am not sure what is going wrong. I have already tried decreasing the learning rate. I started at 0.001 and I am currently at 0.0001