I want to train a neural network in PyTorch, with a single output neuron, to predict the output of a target function specified by me. Let’s say a function of three real variables, so I plan to have a network with 3 input neurons. What is the best way to do this? Should I create a custom dataset with inputs and outputs of the target function, or is there a smarter way? Right now my training loop looks like this (I was experimenting with MNIST on a different network):

for _ in range(training_dict['epochs']):
 for batch, (X, y) in enumerate(dataloader):

  X, y = X.to(device), y.to(device)

  # Compute prediction error
  pred = model(X)

  current_loss = loss(pred, y)

  # Backpropagation
  current_loss.backward()

  optimizer.step()  # Performs a step
  optimizer.zero_grad()
return model

You can see that it works with a dataloader, so again I was wondering if I should create a custom dataset for regression and then use this training loop or if maybe there is a smarter way to do things. In case the standard is to create a custom regression dataset I would like to have an example of how to do that.