Simulation of Binary Black Hole

Hello everyone! I am fairly new to PyTorch and I am currently working on a project that involves using neural networks to describe geodesics around a black hole. I am encountering an error ‘mat1 and mat2 shapes cannot be multiplied (800x3 and 1x32)’ and I am unsure how to resolve it. I would appreciate your help to solve the problem. Please find part of my code below.

# Define the neural network
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(1, 32)
        self.fc2 = nn.Linear(32, 2)

    def forward(self, x):
        x = torch.cos(self.fc1(x))
        x = self.fc2(x)
        return x

# Instantiate the neural network
nn_model = SimpleNN()
# Define the physics-informed neural network (PINNs)
def ODE_model(t, u, nn_model):
    χ, ϕ = u
    p, M, e = ode_model_params

    # Convert input to torch tensor
    input_tensor = torch.tensor([[u[0]]], dtype=torch.float32)

    # Get neural network output and convert to numpy
    with torch.no_grad():
        y = 1 + nn_model(input_tensor).numpy()[0]

    numer = (1 + e * np.cos(χ)) ** 2
    denom = M * (p ** (3 / 2))

    χ_dot = (numer / denom) * y[0]
    ϕ_dot = (numer / denom) * y[1]

    return [χ_dot, ϕ_dot]
# Example parameters
ode_model_params = [100, 1, 0.5]  # [p, M, e]
u0 = [math.pi, 0.0]  
t_span = (0, 60000)  
t_eval = np.linspace(0, 60000, 1000)  

# Solve the ODE
sol_nn = solve_ivp(ODE_model, t_span, u0, args=(nn_model,), t_eval=t_eval, method='RK45')

# Plotting the solution without training (works fine)
# Converting to numpy arrays
real_time = np.array(sol.t)
real_chi = np.array(sol.y[0])
real_phi = np.array(sol.y[1])

nn_time = np.array(sol_nn.t)
nn_chi = np.array(sol_nn.y[0])
nn_phi = np.array(sol_nn.y[1])

X = np.stack((nn_time, nn_chi, nn_phi), axis=-1)
y = np.stack((real_time, real_chi, real_phi), axis=-1)
# Splitting the data
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, train_size = 0.8, test_size = 0.2, random_state = 2)

X_train_tensor = torch.tensor(X_train, dtype=torch.float)
y_train_tensor = torch.tensor(y_train, dtype=torch.float)

X_test_tensor = torch.tensor(X_test, dtype=torch.float)
y_test_tensor = torch.tensor(y_test, dtype=torch.float)
# Loss function and optimization
loss = nn.MSELoss()
optimizer = optim.Adam(nn_model.parameters(), lr = 0.0009)
# Training loop
num_epochs = 10000
train_losses = []
for epochs in range(num_epochs):
    predictions = nn_model(X_train_tensor)
    MSE = loss(predictions, y_train_tensor)
    train_losses.append(MSE.item())
    MSE.backward()
    optimizer.step()
    optimizer.zero_grad()
    if (epochs + 1) % 100 == 0:
        print(f'Epoch [{epochs + 1}/{num_epochs}], MSE Loss: {MSE.item()}')

RuntimeError: mat1 and mat2 shapes cannot be multiplied (800x3 and 1x32)

If you required more information or the files, let me know.
Best regards

Based on the error message it seems the first linear layer raises the shape mismatch. Could you check the input activation shape to make sure it has a shape of [batch_size, 1]?