How can I visualize this NN?

class Net(nn.Module):
def init(self):
super().init()
self.fc1 = nn.Linear(X_train.shape[1], 16)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(16, 1)

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

model = Net()

@ptrblck_de need your help!

Do you want to see like a summary of the layers?

yes and/or input, activation and output nodes with connections if possible

For a summary of the layers with input and outpu you can go for [torchinfo ยท PyPI].

For visualization of the blocks itself, you can go for SummaryWriter and use writer.add_graph.

can you paste the exact lines of code I need here? I keep getting the RuntimeError: mat1 and mat2 shapes cannot be multiplied error

I supposed a random input:

import torch
import torch.nn as nn
from torchinfo import summary


import numpy as np
X_train = np.random.rand(100, 10)  # example: 100 samples, 10 features

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(X_train.shape[1], 16)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(16, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = Net()

# Creation of a  dummy input
batch_size = 32
input_features = X_train.shape[1]
dummy_input = torch.randn(batch_size, input_features)

model_summary = summary(model, 
                        input_size=(batch_size, input_features),
                        col_names=["input_size", "output_size", "num_params", "kernel_size", "mult_adds"],
                        verbose=2)
print(model_summary)