Networking pytorch中的神经网络

我想在PyTorch中创建一个神经网络，它将有2个输入和3个输出以及1个隐藏层。两个输入将是表示图像特征的浮点数，3个输出将是介于0和1之间的实数。例如，输出（1,0,0）表示它是正方形，（0,1,0）表示它是矩形。你知道如何在pytorch中实现吗？

网络的定义如下：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.autograd as autograd 
import torch.nn.functional as F
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self, num_inputs=2, num_outputs=3,hidden_dim=5):
        # define your network here
        super(Net, self).__init__()
        self.layer1 = nn.Linear(num_inputs,hidden_dim)
        self.layer2 = nn.Linear(hidden_dim,num_outputs)


    def forward(self, x):
        # implement the forward pass
        x = F.relu(self.layer1(x))
        x = F.sigmoid(self.layer2(x))
        return x

虽然我已经在这里定义了网络，但您可能应该看看pytorch官方网站上的一些示例，例如如何培训您的模型