As far as I understand your question, you would like to create your model based on different input data.
You could get the number of features from your input and create the model afterwards:
in_features = x_train.shape[1]

model = nn.Sequential(
    nn.Linear(in_features, 2)
)

x = torch.from_numpy(x_train)
output = model(x)